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

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "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,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
    }

"""
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": 100,
        "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,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
    }
    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
        epoch_spectral_loss = 0.0
        t0 = time.time()

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

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

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_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()
            epoch_spectral_loss += spectral_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/spectral_loss", spectral_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"spec={spectral_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
        avg_spectral = epoch_spectral_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} spec={avg_spectral:.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("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    # 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)

+        "pred_weight": args.pred_weight,

+        "input_noise": 0.0,  # will be annealed during training


Now update the training loop to include EMA update and input noise annealing:


Edit

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

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_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()

            # 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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            # 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"]
            spectral_loss = losses["spectral_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()

            # Update EMA target encoder
            model.update_ema()


            # 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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            # 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"]
            spectral_loss = losses["spectral_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()

            # Update EMA target encoder
            model.update_ema()

"""
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": 100,
        "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,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
    }
    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
        epoch_spectral_loss = 0.0
        t0 = time.time()

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

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

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_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()
            epoch_spectral_loss += spectral_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/spectral_loss", spectral_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"spec={spectral_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
        avg_spectral = epoch_spectral_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} spec={avg_spectral:.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("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    # 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)

+            # Input noise annealing: ramp up over first 30% of training

+            # Bridges train/inference gap by simulating imperfect AR inputs

+            noise_progress = min(1.0, global_step / (total_steps * 0.3))

+            model.input_noise = args.input_noise_max * noise_progress

+            # Update EMA target encoder


Now update the argparse defaults:


    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    # 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("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.05)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--free_bits", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    # 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("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.05)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--free_bits", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    # 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)

"""
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": 100,
        "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,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
    }
    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
        epoch_spectral_loss = 0.0
        t0 = time.time()

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

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            # 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"]
            spectral_loss = losses["spectral_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()

            # Update EMA target encoder
            model.update_ema()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            epoch_spectral_loss += spectral_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/spectral_loss", spectral_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"spec={spectral_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
        avg_spectral = epoch_spectral_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} spec={avg_spectral:.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("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    # 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("--kl_weight", type=float, default=0.1)

+    parser.add_argument("--kl_weight", type=float, default=0.05)

+    parser.add_argument("--pred_weight", type=float, default=10.0)

+    parser.add_argument("--ema_decay", type=float, default=0.998)

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

+    parser.add_argument("--input_noise_max", type=float, default=0.2)


Updated task #5 status


Updated task #3 status


Now update inference.py to use the learnable start embedding:


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

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


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

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

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]


        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

-        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale

+        # Use learned start embedding instead of random noise

+        start_emb = self.model.start_emb.to(self.device)


Updated task #4 status


Now let me do the research review — read the final state of all files and verify correctness:


   458→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   459→
   460→class LeWMTTS(nn.Module):
   461→    """
   462→    Complete LeWM TTS model.
   463→
   464→    Training: text + mel → audio embeddings → predict next embedding
   465→    Two losses:
   466→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   467→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   468→    """
   469→
   470→    def __init__(self, config):
   471→        super().__init__()
   472→        d_model = config.get("d_model", 256)
   473→        nhead = config.get("nhead", 4)
   474→        n_mels = config.get("n_mels", 100)
   475→        text_vocab = config.get("text_vocab_size", 256)
   476→        text_layers = config.get("text_encoder_layers", 4)
   477→        audio_layers = config.get("audio_encoder_layers", 4)
   478→        predictor_layers = config.get("predictor_layers", 6)
   479→        dropout = config.get("dropout", 0.1)
   480→        self.kl_weight = config.get("kl_weight", 0.01)
   481→        n_speakers = config.get("n_speakers", 1)
   482→
   483→        self.text_encoder = TextEncoder(
   484→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   485→            num_layers=text_layers, dropout=dropout,
   486→        )
   487→        downsample_factor = config.get("downsample_factor", 4)
   488→        self.audio_encoder = AudioEncoder(
   489→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   490→            num_layers=audio_layers, downsample_factor=downsample_factor,
   491→            dropout=dropout,
   492→        )
   493→        self.predictor = JEPAPredictor(
   494→            d_model=d_model, nhead=nhead,
   495→            num_layers=predictor_layers, dropout=dropout,
   496→        )
   497→
   498→        # Mel decoder: reconstructs mel from latent embeddings
   499→        # Needed so the latent space is grounded to actual audio
   500→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
   501→        self.recon_weight = config.get("recon_weight", 1.0)
   502→
   503→        # EMA target encoder — core JEPA component
   504→        # Provides slowly-moving prediction targets to prevent representation collapse
   505→        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
   506→        for p in self.ema_audio_encoder.parameters():
   507→            p.requires_grad = False
   508→        self.ema_decay = config.get("ema_decay", 0.998)
   509→
   510→        # Speaker embedding (for multi-speaker)
   511→        self.n_speakers = n_speakers
   512→        if n_speakers > 1:
   513→            self.speaker_embed = nn.Embedding(n_speakers, d_model)
   514→        else:
   515→            self.speaker_embed = None
   516→
   517→        # Learnable start-of-sequence embedding for AR inference
   518→        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
   519→
   520→        # Prediction weight — primary JEPA loss
   521→        self.pred_weight = config.get("pred_weight", 10.0)
   522→
   523→        # Input noise injection to bridge train/inference gap
   524→        self.input_noise = config.get("input_noise", 0.0)
   525→
   526→        # Free-bits KL: allow KL up to this threshold without penalty
   527→        self.free_bits = config.get("free_bits", 1.0)
   528→
   529→        # Multi-resolution spectral loss for sharper mel output
   530→        self.spectral_loss = MultiResolutionSpectralLoss()
   531→        self.spectral_weight = config.get("spectral_weight", 0.5)
   532→
   533→        self.config = config
   534→
   535→    def _add_speaker(self, text_emb, speaker_id):
   536→        """Add speaker embedding to text embeddings."""
   537→        if self.speaker_embed is not None and speaker_id is not None:
   538→            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
   539→            text_emb = text_emb + spk
   540→        return text_emb
   541→
   542→    @torch.no_grad()
   543→    def update_ema(self):
   544→        """Update EMA target encoder weights."""
   545→        for p_online, p_ema in zip(self.audio_encoder.parameters(),
   546→                                   self.ema_audio_encoder.parameters()):
   547→            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
   548→
   549→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
   550→        """
   551→        Args:
   552→            mel: [B, n_mels, T_mel]
   553→            text_tokens: [B, T_text]
   554→            mel_mask: [B, T_mel] bool
   555→            text_mask: [B, T_text] bool
   556→        Returns:
   557→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   558→        """
   559→        # Encode text + speaker conditioning
   560→        text_emb = self.text_encoder(text_tokens, text_mask)
   561→        text_emb = self._add_speaker(text_emb, speaker_id)
   562→
   563→        # Online encoder → latent embeddings (used for predictor input + reconstruction)
   564→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   565→
   566→        # EMA target encoder → stable prediction targets (no gradient)
   567→        with torch.no_grad():
   568→            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
   569→            target_emb = mu_ema[:, 1:]  # [B, T-1, d] — EMA targets, shifted
   570→
   571→        # Predictor input: online embeddings, with optional noise injection
   572→        input_emb = z[:, :-1]     # [B, T-1, d]
   573→        if self.training and self.input_noise > 0:
   574→            input_emb = input_emb + torch.randn_like(input_emb) * self.input_noise
   575→
   576→        # Adjust mask for downsampled sequence
   577→        if mel_mask is not None:
   578→            T_down = z.shape[1]
   579→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   580→            pred_mask = ds_mask[:, :-1]
   581→        else:
   582→            pred_mask = None
   583→
   584→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   585→
   586→        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
   587→        if pred_mask is not None:
   588→            valid = ~pred_mask
   589→            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
   590→            prediction_loss = F.mse_loss(
   591→                predicted * valid, target_emb * valid, reduction="sum"
   592→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   593→        else:
   594→            prediction_loss = F.mse_loss(predicted, target_emb)
   595→
   596→        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
   597→        # Prevents both KL collapse (too low) and KL explosion (too high)
   598→        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
   599→        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
   600→        if mel_mask is not None:
   601→            T_down = mu.shape[1]
   602→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   603→            valid_kl = (~ds_mask).unsqueeze(-1).float()
   604→            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
   605→        else:
   606→            kl_loss = kl_free.mean()
   607→
   608→        # ─── Loss 3: Mel reconstruction ───
   609→        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
   610→        T_mel = mel.shape[2]
   611→        T_recon = mel_recon.shape[2]
   612→        if T_recon > T_mel:
   613→            mel_recon = mel_recon[:, :, :T_mel]
   614→        elif T_recon < T_mel:
   615→            mel = mel[:, :, :T_recon]
   616→            if mel_mask is not None:
   617→                mel_mask = mel_mask[:, :T_recon]
   618→
   619→        if mel_mask is not None:
   620→            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
   621→            recon_loss = F.l1_loss(
   622→                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
   623→            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
   624→        else:
   625→            recon_loss = F.l1_loss(mel_recon, mel)
   626→
   627→        # ─── Loss 4: Multi-resolution spectral loss ───
   628→        if mel_mask is not None:
   629→            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
   630→        else:
   631→            spectral_loss = self.spectral_loss(mel_recon, mel)
   632→
   633→        total_loss = (self.pred_weight * prediction_loss
   634→                      + self.kl_weight * kl_loss
   635→                      + self.recon_weight * recon_loss
   636→                      + self.spectral_weight * spectral_loss)
   637→
   638→        return {
   639→            "total_loss": total_loss,
   640→            "prediction_loss": prediction_loss,
   641→            "kl_loss": kl_loss,
   642→            "recon_loss": recon_loss,
   643→            "spectral_loss": spectral_loss,
   644→        }
   645→
   646→    def encode_audio(self, mel):
   647→        """Encode mel to latent embeddings (for inference)."""
   648→        z, mu, _ = self.audio_encoder(mel)
   649→        return mu  # Use mean at inference
   650→
   651→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   652→        """Predict next audio embedding autoregressively."""
   653→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   654→        return predicted[:, -1:]  # Only the last prediction
   655→
   656→    def init_ar_cache(self):
   657→        """Initialize KV cache for autoregressive inference."""
   658→        num_layers = len(self.predictor.transformer.layers)
   659→        return self.predictor.init_cache(num_layers)
   660→
   661→    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   662→        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
   663→        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
   664→
   665→
   666→def count_parameters(model):
   667→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   668→
   669→
   670→def build_model(config=None):
   671→    if config is None:
   672→        config = {
   673→            "d_model": 256,
   674→            "nhead": 4,
   675→            "n_mels": 100,
   676→            "text_vocab_size": 256,
   677→            "text_encoder_layers": 4,
   678→            "audio_encoder_layers": 4,
   679→            "predictor_layers": 6,
   680→            "dropout": 0.1,
   681→            "kl_weight": 0.05,
   682→            "pred_weight": 10.0,
   683→            "ema_decay": 0.998,
   684→            "free_bits": 1.0,
   685→            "input_noise": 0.0,
   686→        }
   687→    model = LeWMTTS(config)

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

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

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

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

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

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

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

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

        # Learnable start-of-sequence embedding for AR inference
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Prediction weight — primary JEPA loss
        self.pred_weight = config.get("pred_weight", 10.0)

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

        # Free-bits KL: allow KL up to this threshold without penalty
        self.free_bits = config.get("free_bits", 1.0)

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

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

    @torch.no_grad()
    def update_ema(self):
        """Update EMA target encoder weights."""
        for p_online, p_ema in zip(self.audio_encoder.parameters(),
                                   self.ema_audio_encoder.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

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

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema[:, 1:]  # [B, T-1, d] — EMA targets, shifted

        # Predictor input: online embeddings, with optional noise injection
        input_emb = z[:, :-1]     # [B, T-1, d]
        if self.training and self.input_noise > 0:
            input_emb = input_emb + torch.randn_like(input_emb) * self.input_noise

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

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

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

        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
        # Prevents both KL collapse (too low) and KL explosion (too high)
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid_kl = (~ds_mask).unsqueeze(-1).float()
            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = kl_free.mean()

        # ─── Loss 3: Mel reconstruction ───
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

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

        # ─── Loss 4: Multi-resolution spectral loss ───
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (self.pred_weight * prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

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

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

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

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

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

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

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.05,
            "pred_weight": 10.0,
            "ema_decay": 0.998,
            "free_bits": 1.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)


     1→"""
     2→Training loop for LeWM TTS.
     3→"""
     4→
     5→import os
     6→import json
     7→import time
     8→import argparse
     9→import torch
    10→from torch.utils.tensorboard import SummaryWriter
    11→from pathlib import Path
    12→
    13→from model import build_model
    14→from dataset import build_dataloader
    15→
    16→
    17→def train(args):
    18→    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    19→    print(f"Device: {device}")
    20→
    21→    # ─── Build model ───
    22→    config = {
    23→        "d_model": args.d_model,
    24→        "nhead": args.nhead,
    25→        "n_mels": 100,
    26→        "text_vocab_size": 256,
    27→        "text_encoder_layers": args.text_layers,
    28→        "audio_encoder_layers": args.audio_layers,
    29→        "predictor_layers": args.predictor_layers,
    30→        "dropout": args.dropout,
    31→        "kl_weight": args.kl_weight,
    32→        "recon_weight": args.recon_weight,
    33→        "spectral_weight": args.spectral_weight,
    34→        "downsample_factor": args.downsample_factor,
    35→        "pred_weight": args.pred_weight,
    36→        "ema_decay": args.ema_decay,
    37→        "free_bits": args.free_bits,
    38→        "input_noise": 0.0,  # will be annealed during training
    39→    }
    40→    model, config = build_model(config)
    41→    model = model.to(device)
    42→
    43→    # ─── Build dataloader ───
    44→    manifest_path = os.path.join(args.data_dir, "manifest.json")
    45→    loader, dataset = build_dataloader(
    46→        manifest_path, batch_size=args.batch_size,
    47→        num_workers=args.num_workers,
    48→    )
    49→    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")
    50→
    51→    # ─── Optimizer & Scheduler ───
    52→    optimizer = torch.optim.AdamW(
    53→        model.parameters(), lr=args.lr,
    54→        betas=(0.9, 0.98), weight_decay=0.01,
    55→    )
    56→
    57→    # Warmup + cosine decay
    58→    total_steps = args.epochs * len(loader)
    59→    warmup_steps = min(4000, total_steps // 10)
    60→
    61→    def lr_lambda(step):
    62→        if step < warmup_steps:
    63→            return step / max(1, warmup_steps)
    64→        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    65→        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())
    66→
    67→    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    68→
    69→    # ─── Logging ───
    70→    ckpt_dir = Path(args.output_dir) / "checkpoints"
    71→    ckpt_dir.mkdir(parents=True, exist_ok=True)
    72→    writer = SummaryWriter(Path(args.output_dir) / "logs")
    73→
    74→    # Save config
    75→    with open(Path(args.output_dir) / "config.json", "w") as f:
    76→        json.dump(config, f, indent=2)
    77→
    78→    # ─── Resume from checkpoint ───
    79→    start_epoch = 0
    80→    global_step = 0
    81→    if args.resume:
    82→        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
    83→        model.load_state_dict(ckpt["model"])
    84→        optimizer.load_state_dict(ckpt["optimizer"])
    85→        scheduler.load_state_dict(ckpt["scheduler"])
    86→        start_epoch = ckpt["epoch"] + 1
    87→        global_step = ckpt["global_step"]
    88→        print(f"Resumed from epoch {start_epoch}, step {global_step}")
    89→
    90→    # ─── Training loop ───
    91→    model.train()
    92→    best_loss = float("inf")
    93→
    94→    for epoch in range(start_epoch, args.epochs):
    95→        epoch_loss = 0.0
    96→        epoch_pred_loss = 0.0
    97→        epoch_kl_loss = 0.0
    98→        epoch_recon_loss = 0.0
    99→        epoch_spectral_loss = 0.0
   100→        t0 = time.time()
   101→
   102→        for batch_idx, batch in enumerate(loader):
   103→            mel = batch["mel"].to(device)
   104→            text_tokens = batch["text_tokens"].to(device)
   105→            mel_mask = batch["mel_mask"].to(device)
   106→            text_mask = batch["text_mask"].to(device)
   107→
   108→            # KL annealing: ramp up over first 20% of training
   109→            kl_anneal = min(1.0, global_step / (total_steps * 0.2))
   110→            model.kl_weight = args.kl_weight * kl_anneal
   111→
   112→            # Input noise annealing: ramp up over first 30% of training
   113→            # Bridges train/inference gap by simulating imperfect AR inputs
   114→            noise_progress = min(1.0, global_step / (total_steps * 0.3))
   115→            model.input_noise = args.input_noise_max * noise_progress
   116→
   117→            # Forward
   118→            losses = model(mel, text_tokens, mel_mask, text_mask)
   119→
   120→            total_loss = losses["total_loss"]
   121→            pred_loss = losses["prediction_loss"]
   122→            kl_loss = losses["kl_loss"]
   123→            recon_loss = losses["recon_loss"]
   124→            spectral_loss = losses["spectral_loss"]
   125→
   126→            # Backward
   127→            optimizer.zero_grad()
   128→            total_loss.backward()
   129→
   130→            # Gradient clipping
   131→            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
   132→
   133→            optimizer.step()
   134→            scheduler.step()
   135→
   136→            # Update EMA target encoder
   137→            model.update_ema()
   138→
   139→            # Logging
   140→            epoch_loss += total_loss.item()
   141→            epoch_pred_loss += pred_loss.item()
   142→            epoch_kl_loss += kl_loss.item()
   143→            epoch_recon_loss += recon_loss.item()
   144→            epoch_spectral_loss += spectral_loss.item()
   145→            global_step += 1
   146→
   147→            if global_step % args.log_every == 0:
   148→                lr = optimizer.param_groups[0]["lr"]
   149→                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
   150→                writer.add_scalar("train/prediction_loss", pred_loss.item(), global_step)
   151→                writer.add_scalar("train/kl_loss", kl_loss.item(), global_step)
   152→                writer.add_scalar("train/recon_loss", recon_loss.item(), global_step)
   153→                writer.add_scalar("train/spectral_loss", spectral_loss.item(), global_step)
   154→                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
   155→                writer.add_scalar("train/lr", lr, global_step)
   156→
   157→            if global_step % args.print_every == 0:
   158→                elapsed = time.time() - t0
   159→                print(
   160→                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
   161→                    f"loss={total_loss.item():.4f} pred={pred_loss.item():.4f} "
   162→                    f"kl={kl_loss.item():.4f} recon={recon_loss.item():.4f} "
   163→                    f"spec={spectral_loss.item():.4f} | "
   164→                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"
   165→                )
   166→
   167→        # Epoch summary
   168→        n_batches = len(loader)
   169→        avg_loss = epoch_loss / n_batches
   170→        avg_pred = epoch_pred_loss / n_batches
   171→        avg_kl = epoch_kl_loss / n_batches
   172→        avg_recon = epoch_recon_loss / n_batches
   173→        avg_spectral = epoch_spectral_loss / n_batches
   174→        epoch_time = time.time() - t0
   175→
   176→        print(
   177→            f"Epoch {epoch+1}/{args.epochs} | "
   178→            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} kl={avg_kl:.4f} "
   179→            f"recon={avg_recon:.4f} spec={avg_spectral:.4f} | {epoch_time:.1f}s"
   180→        )
   181→
   182→        # Save checkpoint
   183→        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
   184→            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
   185→            torch.save({
   186→                "epoch": epoch,
   187→                "global_step": global_step,
   188→                "model": model.state_dict(),
   189→                "optimizer": optimizer.state_dict(),
   190→                "scheduler": scheduler.state_dict(),
   191→                "config": config,
   192→                "loss": avg_loss,
   193→            }, ckpt_path)
   194→            print(f"  Saved checkpoint: {ckpt_path}")
   195→
   196→        # Save best
   197→        if avg_loss < best_loss:
   198→            best_loss = avg_loss
   199→            best_path = ckpt_dir / "best.pt"
   200→            torch.save({
   201→                "epoch": epoch,
   202→                "global_step": global_step,
   203→                "model": model.state_dict(),
   204→                "config": config,
   205→                "loss": avg_loss,
   206→            }, best_path)
   207→            print(f"  New best model (loss={avg_loss:.4f})")
   208→
   209→    writer.close()
   210→    print(f"\nTraining complete. Best loss: {best_loss:.4f}")
   211→
   212→
   213→if __name__ == "__main__":
   214→    parser = argparse.ArgumentParser()
   215→    # Data
   216→    parser.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data")
   217→    parser.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output")
   218→    # Model
   219→    parser.add_argument("--d_model", type=int, default=256)
   220→    parser.add_argument("--nhead", type=int, default=4)
   221→    parser.add_argument("--text_layers", type=int, default=4)
   222→    parser.add_argument("--audio_layers", type=int, default=4)
   223→    parser.add_argument("--predictor_layers", type=int, default=6)
   224→    parser.add_argument("--dropout", type=float, default=0.1)
   225→    parser.add_argument("--downsample_factor", type=int, default=4)
   226→    parser.add_argument("--kl_weight", type=float, default=0.05)
   227→    parser.add_argument("--recon_weight", type=float, default=1.0)
   228→    parser.add_argument("--spectral_weight", type=float, default=0.5)
   229→    parser.add_argument("--pred_weight", type=float, default=10.0)
   230→    parser.add_argument("--ema_decay", type=float, default=0.998)
   231→    parser.add_argument("--free_bits", type=float, default=1.0)
   232→    parser.add_argument("--input_noise_max", type=float, default=0.2)
   233→    # Training
   234→    parser.add_argument("--epochs", type=int, default=200)
   235→    parser.add_argument("--batch_size", type=int, default=16)
   236→    parser.add_argument("--lr", type=float, default=1e-4)
   237→    parser.add_argument("--grad_clip", type=float, default=1.0)
   238→    parser.add_argument("--num_workers", type=int, default=4)
   239→    # Logging
   240→    parser.add_argument("--log_every", type=int, default=10)
   241→    parser.add_argument("--print_every", type=int, default=50)
   242→    parser.add_argument("--save_every", type=int, default=10)
   243→    parser.add_argument("--resume", default=None)
   244→    args = parser.parse_args()
   245→    train(args)
   246→

"""
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": 100,
        "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,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
    }
    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
        epoch_spectral_loss = 0.0
        t0 = time.time()

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

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            # 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"]
            spectral_loss = losses["spectral_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()

            # Update EMA target encoder
            model.update_ema()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            epoch_spectral_loss += spectral_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/spectral_loss", spectral_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"spec={spectral_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
        avg_spectral = epoch_spectral_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} spec={avg_spectral:.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("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.05)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--free_bits", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    # 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)


Let me verify the code compiles and check for issues. Let me also verify the inference code.


Let me verify the code runs without errors and check the inference file:


Bash

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

# Build model with v4 config
config = {
    'd_model': 256, 'nhead': 4, 'n_mels': 100,
    'text_vocab_size': 256,
    'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
    'dropout': 0.1, 'downsample_factor': 4,
    'kl_weight': 0.05, 'recon_weight': 1.0, 'spectral_weight': 0.5,
    'pred_weight': 10.0, 'ema_decay': 0.998, 'free_bits': 1.0, 'input_noise': 0.1,
}
model, config = build_model(config)
model = model.cuda()

# Count trainable params (EMA should NOT be counted)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
ema_params = sum(p.numel() for p in model.ema_audio_encoder.parameters())
print(f'Trainable: {trainable/1e6:.2f}M, Total: {total/1e6:.2f}M, EMA (frozen): {ema_params/1e6:.2f}M')

# Test forward pass
B = 2
mel = torch.randn(B, 100, 200).cuda()
text = torch.randint(0, 256, (B, 50)).cuda()
mel_mask = torch.zeros(B, 200, dtype=torch.bool).cuda()
text_mask = torch.zeros(B, 50, dtype=torch.bool).cuda()

model.train()
losses = model(mel, text, mel_mask, text_mask)
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')

# Test EMA update
model.update_ema()
print('EMA update OK')

# Test start_emb
print(f'start_emb shape: {model.start_emb.shape}, requires_grad: {model.start_emb.requires_grad}')

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

# Verify EMA encoder has no gradients
ema_grads = [p.grad for p in model.ema_audio_encoder.parameters() if p.grad is not None]
print(f'EMA encoder grads: {len(ema_grads)} (should be 0)')
"

Verify model builds and forward/backward work correctly


LeWM TTS model: 24.00M parameters
Trainable: 24.00M, Total: 27.99M, EMA (frozen): 3.99M
  total_loss: 8.1252
  prediction_loss: 0.3998
  kl_loss: 0.0122
  recon_loss: 0.8002
  spectral_loss: 6.6526
EMA update OK
start_emb shape: torch.Size([1, 1, 256]), requires_grad: True
Backward OK
EMA encoder grads: 0 (should be 0)
Shell cwd was reset to /home/ubuntu/pipecat


Everything works. Let me also verify inference works:


    48→    @torch.no_grad()
    49→    def synthesize(self, text, max_steps=300, temperature=0.0):
    50→        """Generate speech from text using AR prediction with KV-cache."""
    51→        text_tokens = self.text_to_tokens(text)
    52→        text_emb = self.model.text_encoder(text_tokens)
    53→
    54→        # Use learned start embedding instead of random noise
    55→        start_emb = self.model.start_emb.to(self.device)
    56→        all_embeddings = [start_emb]
    57→
    58→        cache = self.model.init_ar_cache()
    59→        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
    60→        if temperature > 0:
    61→            next_emb = next_emb + torch.randn_like(next_emb) * temperature
    62→        all_embeddings.append(next_emb)
    63→
    64→        prev_norms = [next_emb.norm().item()]
    65→
    66→        for step in range(1, max_steps):
    67→            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
    68→            if temperature > 0:
    69→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
    70→            all_embeddings.append(next_emb)
    71→
    72→            norm = next_emb.norm().item()
    73→            prev_norms.append(norm)
    74→            if len(prev_norms) > 20:
    75→                recent = prev_norms[-20:]
    76→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
    77→                    break
    78→
    79→        audio_embs = torch.cat(all_embeddings, dim=1)
    80→        mel = self.model.mel_decoder(audio_embs)
    81→        waveform = self.vocos_decode(mel.cpu())
    82→        return waveform.squeeze().numpy(), 24000
    83→
    84→    @torch.no_grad()
    85→    def reconstruct(self, mel_path):
    86→        """Teacher-forced reconstruction: encode real mel → decode back."""
    87→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

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

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

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

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

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


Good. Now let me do the comprehensive research review. Let me also check that the resume mechanism will handle the new EMA state dict properly:


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

config = {
    'd_model': 256, 'nhead': 4, 'n_mels': 100,
    'text_vocab_size': 256,
    'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
    'dropout': 0.1, 'downsample_factor': 4,
    'kl_weight': 0.05, 'recon_weight': 1.0, 'spectral_weight': 0.5,
    'pred_weight': 10.0, 'ema_decay': 0.998, 'free_bits': 1.0, 'input_noise': 0.0,
}
model, config = build_model(config)

# Check state_dict includes EMA
sd = model.state_dict()
ema_keys = [k for k in sd.keys() if 'ema_audio_encoder' in k]
start_keys = [k for k in sd.keys() if 'start_emb' in k]
print(f'EMA keys in state_dict: {len(ema_keys)}')
print(f'start_emb in state_dict: {start_keys}')

# Verify save/load roundtrip
torch.save({'model': sd, 'config': config}, '/tmp/test_ckpt.pt')
model2, _ = build_model(config)
ckpt = torch.load('/tmp/test_ckpt.pt', weights_only=False)
model2.load_state_dict(ckpt['model'])
print('Save/load roundtrip OK')

# Verify loss balance with typical values
# pred=0.05 * 10 = 0.5  (now ~25% of total, up from <1%)
# kl=1.5 * 0.05 = 0.075  (3.75%)
# recon=0.5 * 1.0 = 0.5  (25%)
# spec=0.6 * 0.5 = 0.3   (15%)
# total ≈ 1.375 + whatever free-bits clips
import numpy as np
pred_contrib = 0.05 * 10
kl_contrib = 1.5 * 0.05
recon_contrib = 0.5 * 1.0
spec_contrib = 0.6 * 0.5
total = pred_contrib + kl_contrib + recon_contrib + spec_contrib
print(f'Expected loss contributions:')
print(f'  pred: {pred_contrib:.3f} ({pred_contrib/total*100:.1f}%)')
print(f'  kl:   {kl_contrib:.3f} ({kl_contrib/total*100:.1f}%)')
print(f'  recon:{recon_contrib:.3f} ({recon_contrib/total*100:.1f}%)')
print(f'  spec: {spec_contrib:.3f} ({spec_contrib/total*100:.1f}%)')
print(f'  total:{total:.3f}')
"

Verify checkpoint save/load and loss balance


LeWM TTS model: 24.00M parameters
EMA keys in state_dict: 59
start_emb in state_dict: ['start_emb']
LeWM TTS model: 24.00M parameters
Save/load roundtrip OK
Expected loss contributions:
  pred: 0.500 (36.4%)
  kl:   0.075 (5.5%)
  recon:0.500 (36.4%)
  spec: 0.300 (21.8%)
  total:1.375
Shell cwd was reset to /home/ubuntu/pipecat


The loss balance is much better now. Prediction loss is 36.4% of total (up from <1%), which gives the predictor meaningful training signal. Let me also check the free-bits KL behavior more carefully.

Let me also think about potential issues from a researcher POV:

1. **EMA initialization**: The EMA encoder starts as an exact copy of the online encoder (deepcopy in __init__). This is correct - it means at step 0, the EMA targets are the same as online targets, and they slowly diverge as training progresses.

2. **EMA in state_dict**: The EMA encoder IS saved in model.state_dict() (59 keys). This is correct for checkpoint saving/loading.

3. **Gradient flow**: The EMA encoder parameters don't require gradients and operate within a no_grad block, so backpropagation flows through the predictor and online audio encoder but stops at the EMA encoder—exactly as intended. 4. **Free-bits KL**: With free_bits=1.0, any KL below 1.0 per dimension incurs no penalty, and only the excess gets penalized, which matches the Kingma et al. formulation. 5. **Input noise**: Starts at 0...

Actually, I'm realizing there's a problem with how start_emb is set up. It's defined as a learnable parameter but only used during inference, not in the training forward pass, so it never receives gradient updates. I need to either incorporate it into the training loop somehow—like prepending it to the input sequence—or rethink how to learn it.

If I prepend start_emb to the input during training, the model would learn to predict the first audio embedding from just the start token and text context, while subsequent positions predict normally from the previous embedding. This is actually crucial for inference because the first step needs to know how to begin generating audio. I'm working through the dimensions here—with the prepended start embedding, the input sequence becomes T_down elements long, which matches the target z_ema length perfectly for a clean 1-1 mapping. I also need to handle the mask carefully by adding an unmasked position at the start for the learnable embedding, then prepending it to the input sequence.

Modifying the forward method... I'm realizing I need to be more careful about how the masks align with the targets and inputs—the target mask should match the original downsampled mask since padded positions in z remain padded in z_ema, while the input mask needs the same structure but with position 0 always unmasked since that's the start token. Now I'm setting up the masking logic where the input mask treats the start embedding as always valid, then shifts the downsampled mask back by one position for predictions. The target mask directly follows the downsampled mask. Then I'm computing the MSE loss while masking out invalid positions...

Wait, I need to clarify the distinction here—the pred_mask serves as the key padding mask for the predictor's self-attention layer, telling it which positions to ignore, while the target_mask is specifically for determining which output positions should contribute to the loss calculation. These serve different purposes in the forward pass versus the loss computation.

For the input side, position 0 is always the start embedding (valid), and positions after that are valid based on whether the corresponding downsampled mask position is False. For the target side, validity depends directly on the downsampled mask. So I need to construct these masks separately—prepending the start embedding to the input sequence and shifting the downsampled mask appropriately for the predictor, while using the unshifted mask for loss computation.

Building the corrected masking logic... Finishing the mask concatenation and setting up the loss computation with proper masking for valid positions, then running the predictor and calculating the mean squared error loss that accounts for which tokens should contribute to the gradient.

Now I need to reconsider how noise is applied during training—the start embedding should remain clean since inference uses it directly, but subsequent positions should be noised to match the imperfect predicted embeddings that would be fed back at test time. Let me adjust the forward method to handle this distinction.