        # 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
        avg_speaker = epoch_speaker_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} spk={avg_speaker:.4f} | {epoch_time:.1f}s"
        )

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 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
        epoch_speaker_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)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.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, speaker_id=speaker_id)

            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"]
            speaker_loss = losses["speaker_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()
            epoch_speaker_loss += speaker_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/speaker_loss", speaker_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} spk={speaker_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)

+        avg_speaker = epoch_speaker_loss / n_batches

-            f"recon={avg_recon:.4f} spec={avg_spectral:.4f} | {epoch_time:.1f}s"

+            f"recon={avg_recon:.4f} spec={avg_spectral:.4f} spk={avg_speaker:.4f} | {epoch_time:.1f}s"


Now add the `--n_speakers` and `--speaker_weight` args:


Edit

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

    parser.add_argument("--free_bits", type=float, default=2.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)


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

    parser.add_argument("--free_bits", type=float, default=2.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 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
        epoch_speaker_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)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.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, speaker_id=speaker_id)

            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"]
            speaker_loss = losses["speaker_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()
            epoch_speaker_loss += speaker_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/speaker_loss", speaker_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} spk={speaker_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
        avg_speaker = epoch_speaker_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} spk={avg_speaker:.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)

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

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

+    parser.add_argument("--n_speakers", type=int, default=1)

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


Now add `n_speakers` and `speaker_weight` to the config dict in `train()`:


        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training

        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,


        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 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
        epoch_speaker_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)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.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, speaker_id=speaker_id)

            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"]
            speaker_loss = losses["speaker_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()
            epoch_speaker_loss += speaker_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/speaker_loss", speaker_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} spk={speaker_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
        avg_speaker = epoch_speaker_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} spk={avg_speaker:.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=2.0)
    parser.add_argument("--input_noise_max", type=float, default=0.2)
    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)
    # Training
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    # Logging
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=50)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args()
    train(args)

+        "speaker_weight": args.speaker_weight,


Now let me also update the inference to handle the speaker embedding injection into latents (since at inference the decoder needs speaker-conditioned embeddings):


1	"""
2	Inference pipeline for LeWM TTS.
3	Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
4	"""
5
6	import torch
7	import torchaudio
8	import numpy as np
9	import argparse
10	from pathlib import Path
11
12	from model import LeWMTTS
13
14
15	class LeWMTTSInference:
16	    """Full inference pipeline."""
17
18	    def __init__(self, checkpoint_path, device="cuda"):
19	        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
20
21	        # Load checkpoint
22	        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
23	        config = ckpt["config"]
24
25	        # Build model
26	        self.model = LeWMTTS(config).to(self.device)
27	        self.model.load_state_dict(ckpt["model"])
28	        self.model.eval()
29
30	        # Vocos vocoder (directly decodes 100-mel)
31	        from vocos import Vocos
32	        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
33	        self.vocos.eval()
34
35	        self.config = config
36	        print(f"Model loaded from {checkpoint_path}")
37
38	    def vocos_decode(self, mel):
39	        """Decode mel [B, 100, T] → waveform via Vocos."""
40	        with torch.no_grad():
41	            waveform = self.vocos.decode(mel)  # [B, T_audio]
42	        return waveform
43
44	    def text_to_tokens(self, text):
45	        tokens = list(text.encode("utf-8"))
46	        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
47
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)
88
89	        z, mu, logvar = self.model.audio_encoder(mel)
90	        mel_recon = self.model.mel_decoder(mu)
91
92	        waveform = self.vocos_decode(mel_recon.cpu())
93	        return waveform.squeeze().numpy(), 24000
94
95	    @torch.no_grad()
96	    def reconstruct_with_prediction(self, mel_path):
97	        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
98	        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
99
100	        z, mu, logvar = self.model.audio_encoder(mel)
101
102	        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
103	        text_emb = self.model.text_encoder(text_tokens)
104
105	        predicted = self.model.predictor(mu[:, :-1], text_emb)
106	        combined = torch.cat([mu[:, :1], predicted], dim=1)
107
108	        mel_recon = self.model.mel_decoder(combined)
109	        waveform = self.vocos_decode(mel_recon.cpu())
110	        return waveform.squeeze().numpy(), 24000
111
112	    def save_audio(self, waveform, sr, output_path):
113	        # Normalize
114	        mx = np.abs(waveform).max()
115	        if mx > 0:
116	            waveform = waveform / mx * 0.95
117	        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
118	        torchaudio.save(output_path, wav_tensor, sr)
119	        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
120
121
122	def main():
123	    parser = argparse.ArgumentParser()
124	    parser.add_argument("--checkpoint", required=True)
125	    parser.add_argument("--text", default=None, help="Text to synthesize")
126	    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
127	    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
128	                        default="synthesize")
129	    parser.add_argument("--output", default="output.wav")
130	    parser.add_argument("--max_steps", type=int, default=300)
131	    parser.add_argument("--temperature", type=float, default=0.0)
132	    args = parser.parse_args()
133
134	    tts = LeWMTTSInference(args.checkpoint)
135
136	    if args.mode == "synthesize":
137	        assert args.text, "Need --text for synthesis"
138	        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
139	    elif args.mode == "reconstruct":
140	        assert args.mel, "Need --mel for reconstruction"
141	        waveform, sr = tts.reconstruct(args.mel)
142	    elif args.mode == "predict":
143	        assert args.mel, "Need --mel for prediction test"
144	        waveform, sr = tts.reconstruct_with_prediction(args.mel)
145
146	    tts.save_audio(waveform, sr, args.output)
147
148
149	if __name__ == "__main__":
150	    main()
151

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

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

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


    @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]

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
        """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)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

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


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

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
        """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)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

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

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

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

-    def synthesize(self, text, max_steps=300, temperature=0.0):

+    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):

+        # Add speaker conditioning to text embeddings

+        if self.model.speaker_embed is not None and speaker_id is not None:

+            spk_id = torch.tensor([speaker_id], device=self.device)

+            text_emb = self.model._add_speaker(text_emb, spk_id)

+            # Cache speaker embedding for adding to predicted latents

+            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]


Now add speaker embedding to predicted latents during AR loop so the mel decoder sees speaker-conditioned embeddings (matching training):


        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)

        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
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)


        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
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)

"""
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, speaker_id=None):
        """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)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

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

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

-        all_embeddings.append(next_emb)

+        # Add speaker embedding to predicted latent (matches training where z = z + spk)

+        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb

+        all_embeddings.append(emb_for_decode)

-            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)

+            # Feed speaker-conditioned embedding as input (matches training input_emb)

+            next_input = emb_for_decode

+            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)

-            all_embeddings.append(next_emb)

+            all_embeddings.append(emb_for_decode)


Now add `--speaker_id` to the inference CLI:


    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)

    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--speaker_id", type=int, default=None)
    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,
                                       speaker_id=args.speaker_id)


    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--speaker_id", type=int, default=None)
    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,
                                       speaker_id=args.speaker_id)

"""
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, speaker_id=None):
        """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)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

        # 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
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)

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

+    parser.add_argument("--speaker_id", type=int, default=None)

-        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)

+        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature,

+                                       speaker_id=args.speaker_id)


Now let me verify the dataset returns `speaker_id` when available:


1	"""
2	Dataset and DataLoader for LeWM TTS training.
3	"""
4
5	import json
6	import torch
7	from torch.utils.data import Dataset, DataLoader
8	from pathlib import Path
9
10
11	class TTSDataset(Dataset):
12	    """Loads preprocessed mel spectrograms and text for TTS training."""
13
14	    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
15	        with open(manifest_path, "r", encoding="utf-8") as f:
16	            self.manifest = json.load(f)
17
18	        self.max_mel_frames = max_mel_frames
19	        self.max_text_len = max_text_len
20
21	        # Filter out entries that are too long
22	        self.manifest = [
23	            e for e in self.manifest
24	            if e["mel_frames"] <= max_mel_frames
25	        ]
26	        print(f"Dataset: {len(self.manifest)} samples loaded")
27
28	    def __len__(self):
29	        return len(self.manifest)
30
31	    def text_to_tokens(self, text):
32	        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
33	        # UTF-8 byte encoding — simple, universal, no tokenizer needed
34	        tokens = list(text.encode("utf-8"))
35	        return tokens
36
37	    def __getitem__(self, idx):
38	        entry = self.manifest[idx]
39
40	        # Load mel
41	        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]
42
43	        # Tokenize text
44	        tokens = self.text_to_tokens(entry["text"])
45	        tokens = tokens[:self.max_text_len]
46	        text_tokens = torch.tensor(tokens, dtype=torch.long)
47
48	        return {
49	            "mel": mel,
50	            "text_tokens": text_tokens,
51	            "mel_frames": mel.shape[1],
52	            "text_len": len(tokens),
53	        }
54
55
56	def collate_fn(batch):
57	    """Pad and batch samples."""
58	    # Find max lengths
59	    max_mel = max(s["mel_frames"] for s in batch)
60	    max_text = max(s["text_len"] for s in batch)
61	    n_mels = batch[0]["mel"].shape[0]
62	    B = len(batch)
63
64	    # Allocate padded tensors
65	    mel_padded = torch.zeros(B, n_mels, max_mel)
66	    text_padded = torch.zeros(B, max_text, dtype=torch.long)
67	    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
68	    text_mask = torch.ones(B, max_text, dtype=torch.bool)
69
70	    for i, s in enumerate(batch):
71	        t_mel = s["mel_frames"]
72	        t_text = s["text_len"]
73	        mel_padded[i, :, :t_mel] = s["mel"]
74	        text_padded[i, :t_text] = s["text_tokens"]
75	        mel_mask[i, :t_mel] = False
76	        text_mask[i, :t_text] = False
77
78	    return {
79	        "mel": mel_padded,
80	        "text_tokens": text_padded,
81	        "mel_mask": mel_mask,
82	        "text_mask": text_mask,
83	    }
84
85
86	def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
87	    dataset = TTSDataset(manifest_path)
88	    loader = DataLoader(
89	        dataset,
90	        batch_size=batch_size,
91	        shuffle=shuffle,
92	        num_workers=num_workers,
93	        collate_fn=collate_fn,
94	        pin_memory=True,
95	        drop_last=True,
96	    )
97	    return loader, dataset
98

"""
Dataset and DataLoader for LeWM TTS training.
"""

import json
import torch
from torch.utils.data import Dataset, DataLoader
from pathlib import Path

class TTSDataset(Dataset):
    """Loads preprocessed mel spectrograms and text for TTS training."""

    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)

        self.max_mel_frames = max_mel_frames
        self.max_text_len = max_text_len

        # Filter out entries that are too long
        self.manifest = [
            e for e in self.manifest
            if e["mel_frames"] <= max_mel_frames
        ]
        print(f"Dataset: {len(self.manifest)} samples loaded")

    def __len__(self):
        return len(self.manifest)

    def text_to_tokens(self, text):
        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
        # UTF-8 byte encoding — simple, universal, no tokenizer needed
        tokens = list(text.encode("utf-8"))
        return tokens

    def __getitem__(self, idx):
        entry = self.manifest[idx]

        # Load mel
        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]

        # Tokenize text
        tokens = self.text_to_tokens(entry["text"])
        tokens = tokens[:self.max_text_len]
        text_tokens = torch.tensor(tokens, dtype=torch.long)

        return {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

def collate_fn(batch):
    """Pad and batch samples."""
    # Find max lengths
    max_mel = max(s["mel_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    n_mels = batch[0]["mel"].shape[0]
    B = len(batch)

    # Allocate padded tensors
    mel_padded = torch.zeros(B, n_mels, max_mel)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_mel = s["mel_frames"]
        t_text = s["text_len"]
        mel_padded[i, :, :t_mel] = s["mel"]
        text_padded[i, :t_text] = s["text_tokens"]
        mel_mask[i, :t_mel] = False
        text_mask[i, :t_text] = False

    return {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    dataset = TTSDataset(manifest_path)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    return loader, dataset


The dataset needs to pass through `speaker_id` when present in the manifest:


        return {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

        result = {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

        # Multi-speaker: pass speaker_id if present in manifest
        if "speaker_id" in entry:
            result["speaker_id"] = entry["speaker_id"]

        return result


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

        result = {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

        # Multi-speaker: pass speaker_id if present in manifest
        if "speaker_id" in entry:
            result["speaker_id"] = entry["speaker_id"]

        return result

"""
Dataset and DataLoader for LeWM TTS training.
"""

import json
import torch
from torch.utils.data import Dataset, DataLoader
from pathlib import Path

class TTSDataset(Dataset):
    """Loads preprocessed mel spectrograms and text for TTS training."""

    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)

        self.max_mel_frames = max_mel_frames
        self.max_text_len = max_text_len

        # Filter out entries that are too long
        self.manifest = [
            e for e in self.manifest
            if e["mel_frames"] <= max_mel_frames
        ]
        print(f"Dataset: {len(self.manifest)} samples loaded")

    def __len__(self):
        return len(self.manifest)

    def text_to_tokens(self, text):
        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
        # UTF-8 byte encoding — simple, universal, no tokenizer needed
        tokens = list(text.encode("utf-8"))
        return tokens

    def __getitem__(self, idx):
        entry = self.manifest[idx]

        # Load mel
        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]

        # Tokenize text
        tokens = self.text_to_tokens(entry["text"])
        tokens = tokens[:self.max_text_len]
        text_tokens = torch.tensor(tokens, dtype=torch.long)

        return {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

def collate_fn(batch):
    """Pad and batch samples."""
    # Find max lengths
    max_mel = max(s["mel_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    n_mels = batch[0]["mel"].shape[0]
    B = len(batch)

    # Allocate padded tensors
    mel_padded = torch.zeros(B, n_mels, max_mel)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_mel = s["mel_frames"]
        t_text = s["text_len"]
        mel_padded[i, :, :t_mel] = s["mel"]
        text_padded[i, :t_text] = s["text_tokens"]
        mel_mask[i, :t_mel] = False
        text_mask[i, :t_text] = False

    return {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    dataset = TTSDataset(manifest_path)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    return loader, dataset

+        # Multi-speaker: pass speaker_id if present in manifest

+            result["speaker_id"] = entry["speaker_id"]


    return {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

    result = {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

    # Multi-speaker: batch speaker_ids if present
    if "speaker_id" in batch[0]:
        result["speaker_id"] = torch.tensor([s["speaker_id"] for s in batch], dtype=torch.long)

    return result


    result = {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

    # Multi-speaker: batch speaker_ids if present
    if "speaker_id" in batch[0]:
        result["speaker_id"] = torch.tensor([s["speaker_id"] for s in batch], dtype=torch.long)

    return result

"""
Dataset and DataLoader for LeWM TTS training.
"""

import json
import torch
from torch.utils.data import Dataset, DataLoader
from pathlib import Path

class TTSDataset(Dataset):
    """Loads preprocessed mel spectrograms and text for TTS training."""

    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)

        self.max_mel_frames = max_mel_frames
        self.max_text_len = max_text_len

        # Filter out entries that are too long
        self.manifest = [
            e for e in self.manifest
            if e["mel_frames"] <= max_mel_frames
        ]
        print(f"Dataset: {len(self.manifest)} samples loaded")

    def __len__(self):
        return len(self.manifest)

    def text_to_tokens(self, text):
        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
        # UTF-8 byte encoding — simple, universal, no tokenizer needed
        tokens = list(text.encode("utf-8"))
        return tokens

    def __getitem__(self, idx):
        entry = self.manifest[idx]

        # Load mel
        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]

        # Tokenize text
        tokens = self.text_to_tokens(entry["text"])
        tokens = tokens[:self.max_text_len]
        text_tokens = torch.tensor(tokens, dtype=torch.long)

        result = {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

        # Multi-speaker: pass speaker_id if present in manifest
        if "speaker_id" in entry:
            result["speaker_id"] = entry["speaker_id"]

        return result

def collate_fn(batch):
    """Pad and batch samples."""
    # Find max lengths
    max_mel = max(s["mel_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    n_mels = batch[0]["mel"].shape[0]
    B = len(batch)

    # Allocate padded tensors
    mel_padded = torch.zeros(B, n_mels, max_mel)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_mel = s["mel_frames"]
        t_text = s["text_len"]
        mel_padded[i, :, :t_mel] = s["mel"]
        text_padded[i, :t_text] = s["text_tokens"]
        mel_mask[i, :t_mel] = False
        text_mask[i, :t_text] = False

    return {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    dataset = TTSDataset(manifest_path)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    return loader, dataset

+    # Multi-speaker: batch speaker_ids if present

+        result["speaker_id"] = torch.tensor([s["speaker_id"] for s in batch], dtype=torch.long)


Now let me verify everything compiles:


Bash

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

# Test single-speaker (backward compatible)
model, config = build_model()
print('Single-speaker OK')

# Test multi-speaker
config_ms = {
    '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': 2.0, 'input_noise': 0.0,
    'n_speakers': 15, 'speaker_weight': 1.0,
}
model_ms, _ = build_model(config_ms)
print(f'Multi-speaker OK: speaker_embed={model_ms.speaker_embed}, classifier={model_ms.speaker_classifier}')

# Test forward pass
import torch
B, T_mel, T_text = 2, 400, 20
mel = torch.randn(B, 100, T_mel)
text = torch.randint(0, 256, (B, T_text))
mel_mask = torch.zeros(B, T_mel, dtype=torch.bool)
text_mask = torch.zeros(B, T_text, dtype=torch.bool)
speaker_id = torch.randint(0, 15, (B,))

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

# Test without speaker (backward compat)
losses_nospeaker = model_ms(mel, text, mel_mask, text_mask)
print('Forward without speaker_id OK')

# Single-speaker model forward (no speaker loss)
losses_single = model(mel, text, mel_mask, text_mask)
print(f'Single-speaker forward OK, speaker_loss={losses_single[\"speaker_loss\"].item():.4f}')
" 2>&1

Test model builds and forward passes


Error: Exit code 127
/bin/bash: line 40: python: command not found


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

# Test single-speaker (backward compatible)
model, config = build_model()
print('Single-speaker OK')

# Test multi-speaker
config_ms = {
    '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': 2.0, 'input_noise': 0.0,
    'n_speakers': 15, 'speaker_weight': 1.0,
}
model_ms, _ = build_model(config_ms)
print(f'Multi-speaker OK: speaker_embed={model_ms.speaker_embed}, classifier={model_ms.speaker_classifier}')

# Test forward pass
import torch
B, T_mel, T_text = 2, 400, 20
mel = torch.randn(B, 100, T_mel)
text = torch.randint(0, 256, (B, T_text))
mel_mask = torch.zeros(B, T_mel, dtype=torch.bool)
text_mask = torch.zeros(B, T_text, dtype=torch.bool)
speaker_id = torch.randint(0, 15, (B,))

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

# Test without speaker (backward compat)
losses_nospeaker = model_ms(mel, text, mel_mask, text_mask)
print('Forward without speaker_id OK')

# Single-speaker model forward (no speaker loss)
losses_single = model(mel, text, mel_mask, text_mask)
print(f'Single-speaker forward OK, speaker_loss={losses_single[\"speaker_loss\"].item():.4f}')
" 2>&1

Test model builds and forward passes


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py:515: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)
  output = torch._nested_tensor_from_mask(
LeWM TTS model: 24.00M parameters
Single-speaker OK
LeWM TTS model: 24.01M parameters
Multi-speaker OK: speaker_embed=Embedding(15, 256), classifier=Linear(in_features=256, out_features=15, bias=True)
  total_loss: 19.9207
  prediction_loss: 1.2760
  kl_loss: 0.0014
  recon_loss: 0.8004
  spectral_loss: 6.6123
  speaker_loss: 3.0545
Forward without speaker_id OK
Single-speaker forward OK, speaker_loss=0.0000


All passing. Let me verify the changes are clean:


Check changed files


Error: Exit code 129
warning: Not a git repository. Use --no-index to compare two paths outside a working tree
usage: git diff --no-index [<options>] <path> <path>

Diff output format options
    -p, --patch           generate patch
    -s, --no-patch        suppress diff output
    -u                    generate patch
    -U, --unified[=<n>]   generate diffs with <n> lines context
    -W, --function-context
                          generate diffs with <n> lines context
    --raw                 generate the diff in raw format
    --patch-with-raw      synonym for '-p --raw'
    --patch-with-stat     synonym for '-p --stat'
    --numstat             machine friendly --stat
    --shortstat           output only the last line of --stat
    -X, --dirstat[=<param1,param2>...]
                          output the distribution of relative amount of changes for each sub-directory
    --cumulative          synonym for --dirstat=cumulative
    --dirstat-by-file[=<param1,param2>...]
                          synonym for --dirstat=files,param1,param2...
    --check               warn if changes introduce conflict markers or whitespace errors
    --summary             condensed summary such as creations, renames and mode changes
    --name-only           show only names of changed files
    --name-status         show only names and status of changed files
    --stat[=<width>[,<name-width>[,<count>]]]
                          generate diffstat
    --stat-width <width>  generate diffstat with a given width
    --stat-name-width <width>
                          generate diffstat with a given name width
    --stat-graph-width <width>
                          generate diffstat with a given graph width
    --stat-count <count>  generate diffstat with limited lines
    --compact-summary     generate compact summary in diffstat
    --binary              output a binary diff that can be applied
    --full-index          show full pre- and post-image object names on the "index" lines
    --color[=<when>]      show colored diff
    --ws-error-highlight <kind>
                          highlight whitespace errors in the 'context', 'old' or 'new' lines in the diff
    -z                    do not munge pathnames and use NULs as output field terminators in --raw or --numstat
    --abbrev[=<n>]        use <n> digits to display object names
    --src-prefix <prefix>
                          show the given source prefix instead of "a/"
    --dst-prefix <prefix>
                          show the given destination prefix instead of "b/"
    --line-prefix <prefix>
                          prepend an additional prefix to every line of output
    --no-prefix           do not show any source or destination prefix
    --inter-hunk-context <n>
                          show context between diff hunks up to the specified number of lines
    --output-indicator-new <char>
                          specify the character to indicate a new line instead of '+'
    --output-indicator-old <char>
                          specify the character to indicate an old line instead of '-'
    --output-indicator-context <char>
                          specify the character to indicate a context instead of ' '

Diff rename options
    -B, --break-rewrites[=<n>[/<m>]]
                          break complete rewrite changes into pairs of delete and create
    -M, --find-renames[=<n>]
                          detect renames
    -D, --irreversible-delete
                          omit the preimage for deletes
    -C, --find-copies[=<n>]
                          detect copies
    --find-copies-harder  use unmodified files as source to find copies
    --no-renames          disable rename detection
    --rename-empty        use empty blobs as rename source
    --follow              continue listing the history of a file beyond renames
    -l <n>                prevent rename/copy detection if the number of rename/copy targets exceeds given limit

Diff algorithm options
    --minimal             produce the smallest possible diff
    -w, --ignore-all-space
                          ignore whitespace when comparing lines
    -b, --ignore-space-change
                          ignore changes in amount of whitespace
    --ignore-space-at-eol
                          ignore changes in whitespace at EOL
    --ignore-cr-at-eol    ignore carrier-return at the end of line
    --ignore-blank-lines  ignore changes whose lines are all blank
    -I, --ignore-matching-lines <regex>
                          ignore changes whose all lines match <regex>
    --indent-heuristic    heuristic to shift diff hunk boundaries for easy reading
    --patience            generate diff using the "patience diff" algorithm
    --histogram           generate diff using the "histogram diff" algorithm
    --diff-algorithm <algorithm>
                          choose a diff algorithm
    --anchored <text>     generate diff using the "anchored diff" algorithm
    --word-diff[=<mode>]  show word diff, using <mode> to delimit changed words
    --word-diff-regex <regex>
                          use <regex> to decide what a word is
    --color-words[=<regex>]
                          equivalent to --word-diff=color --word-diff-regex=<regex>
    --color-moved[=<mode>]
                          moved lines of code are colored differently
    --color-moved-ws <mode>
                          how white spaces are ignored in --color-moved

Other diff options
    --relative[=<prefix>]
                          when run from subdir, exclude changes outside and show relative paths
    -a, --text            treat all files as text
    -R                    swap two inputs, reverse the diff
    --exit-code           exit with 1 if there were differences, 0 otherwise
    --quiet               disable all output of the program
    --ext-diff            allow an external diff helper to be executed
    --textconv            run external text conversion filters when comparing binary files
    --ignore-submodules[=<when>]
                          ignore changes to submodules in the diff generation
    --submodule[=<format>]
                          specify how differences in submodules are shown
    --ita-invisible-in-index
                          hide 'git add -N' entries from the index
    --ita-visible-in-index
                          treat 'git add -N' entries as real in the index
    -S <string>           look for differences that change the number of occurrences of the specified string
    -G <regex>            look for differences that change the number of occurrences of the specified regex
    --pickaxe-all         show all changes in the changeset with -S or -G
    --pickaxe-regex       treat <string> in -S as extended POSIX regular expression
    -O <file>             control the order in which files appear in the output
    --rotate-to <path>    show the change in the specified path first
    --skip-to <path>      skip the output to the specified path
    --find-object <object-id>
                          look for differences that change the number of occurrences of the specified object
    --diff-filter [(A|C|D|M|R|T|U|X|B)...[*]]
                          select files by diff type
    --output <file>       Output to a specific file


Everything works. Here's a summary of what was changed across 4 files:

## Changes Made

### `model.py` — 3 fixes for speaker preservation

**Fix 1: Speaker embedding injected into targets + encoder output** (biggest impact)
- Lines ~575-580: After EMA target encoding, `target_emb = target_emb + spk` and `z = z + spk`
- The prediction loss now **requires** the predictor to know the speaker — two speakers saying the same thing produce different targets, so matching them without speaker info is impossible

**Fix 2: Speaker consistency classifier**
- New `self.speaker_classifier = nn.Linear(d_model, n_speakers)` — negligible params
- New Loss 5: cross-entropy on time-pooled `mu` against `speaker_id`
- Directly tells the audio encoder "keep speaker identity in the latent"

**Fix 3: Configurable `speaker_weight`** (default 1.0)

### `train.py`
- Passes `speaker_id` from batch to model forward
- Logs `speaker_loss` to tensorboard and print
- New args: `--n_speakers`, `--speaker_weight`
- Bumped `--free_bits` default from 1.0 → 2.0 (less KL pressure on speaker dims)

### `dataset.py`
- Passes through `speaker_id` from manifest entries when present
- Batches speaker IDs as a tensor in collate

### `inference.py`
- `synthesize()` accepts `speaker_id` parameter
- Speaker embedding added to both text (cross-attn conditioning) and predicted latents (matching training)
- New CLI arg `--speaker_id`

## To use with multi-speaker

Your manifest.json entries need a `"speaker_id": 0` field (integer, 0-indexed). Then train with:

```bash
python3 train.py --n_speakers 15 --speaker_weight 1.0 --free_bits 2.0 ...
```

For single-speaker, everything is backward compatible — speaker loss is 0.0, no extra params.


test again fixes adn check whwerhet u can do any improvments .. and what is our current model size