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

"""
Train predictor + text encoder + duration predictor with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
Text is expanded to frame level via duration predictor — can't be ignored.
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import (
    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,
    DurationPredictor, compute_uniform_durations, length_regulate,
)
from dataset_v5 import build_dataloader

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

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["dur_weight"] = args.dur_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)

    # Load what we can from pretrained (projections, etc.)
    pretrained_state = ckpt["model"]
    model_state = model.state_dict()
    loaded_keys = []
    skipped_keys = []
    for k, v in pretrained_state.items():
        if k in model_state and model_state[k].shape == v.shape:
            model_state[k] = v
            loaded_keys.append(k)
        else:
            skipped_keys.append(k)
    model.load_state_dict(model_state)
    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")
    if skipped_keys:
        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

    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, max_codec_frames=args.max_codec_frames,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(2000, 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)

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

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

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

            # Text encoding (trainable)
            text_emb = model.text_encoder(text_tokens, text_mask)

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

            # Prepend start, shift right
            start = model.start_emb.expand(B, -1, -1)
            input_emb = torch.cat([start, z[:, :-1]], dim=1)

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

        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: {ckpt_path}")

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)

-Train ONLY the predictor + text encoder with frozen projections.

+Train predictor + text encoder + duration predictor with frozen projections.

-No EMA — targets come directly from frozen proj_in (stable, no divergence).

+Text is expanded to frame level via duration predictor — can't be ignored.

-from model_v5 import LeWMTTSv5, build_model

+    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,

+    DurationPredictor, compute_uniform_durations, length_regulate,

+    config["dur_weight"] = args.dur_weight

-    model.load_state_dict(ckpt["model"])

-    print(f"Loaded pretrained from {args.pretrained} (epoch {ckpt['epoch']+1})")

+    # Load what we can from pretrained (projections, etc.)

+    for k, v in pretrained_state.items():

+        if k in model_state and model_state[k].shape == v.shape:

+    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")

+        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

for p in model.ema_proj_in.parameters():

-    # Reinitialize predictor + text encoder from scratch if requested

+    # Reinitialize predictor + text encoder + duration predictor from scratch

-        print("Reinitializing predictor + text encoder from scratch")

-        from model_v5 import JEPAPredictor, TextEncoder

+        print("Reinitializing predictor + text encoder + duration predictor from scratch")

+        model.duration_predictor = DurationPredictor(

+            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,

-    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + start_emb)")

+    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

for batch_idx, batch in enumerate(loader):

+            log_dur_pred = model.duration_predictor(text_emb, text_mask)

+            # Ground truth durations (uniform)

+            text_lengths = (~text_mask).sum(dim=1)

+            audio_lengths = (~codec_mask).sum(dim=1)

+            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)

+            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

+            valid_text = (~text_mask).float()

+                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"

+            ) / (valid_text.sum() + 1e-8)

+            # Expand text to frame level using GT durations

+            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)

+            if text_emb_expanded.shape[1] > T:

+                text_emb_expanded = text_emb_expanded[:, :T]

+            elif text_emb_expanded.shape[1] < T:

+                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],

+                                device=device)

+                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

z = model.proj_in(codec_seq)  # [B, T, d_model]

-                target_emb = z.clone()  # no EMA, direct target

start = model.start_emb.expand(B, -1, -1)

input_emb = torch.cat([start, z[:, :-1]], dim=1)

-            # Classifier-free guidance: randomly drop audio context

-            # Forces the model to use text cross-attention

-            if args.audio_drop_prob > 0:

-                drop_mask = torch.rand(B, device=device) < args.audio_drop_prob

-                    input_emb = input_emb.clone()

-                    input_emb[drop_mask] = 0.0  # zero out audio for dropped samples

-                start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)

-                pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

+            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)

+            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

-            predicted = model.predictor(input_emb, text_emb, pred_mask, text_mask)

+            # Predict with frame-level text

+            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

# Prediction loss (MSE vs frozen proj_in targets)

-                valid = (~loss_mask).unsqueeze(-1)

-                pred_loss = F.mse_loss(

-                    predicted * valid, target_emb * valid, reduction="sum"

-                ) / (valid.sum() * predicted.shape[-1] + 1e-8)

-                pred_loss = F.mse_loss(predicted, target_emb)

+            valid = (~codec_mask).unsqueeze(-1)

+                predicted * valid, target_emb * valid, reduction="sum"

+            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

+            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

+            epoch_pred += pred_loss.item()

+            epoch_dur += dur_loss.item()

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

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

-                    f"pred={pred_loss.item():.4f} | "

+                    f"pred={pred_loss.item():.4f} dur={dur_loss.item():.4f} | "

+        avg_dur = epoch_dur / len(loader)

-        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} | {epoch_time:.1f}s")

+        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

-            # Save full model (including frozen projections)

ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"

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

-    parser.add_argument("--audio_drop_prob", type=float, default=0.0, help="Prob of dropping audio context (CFG)")

"""
Train ONLY the predictor + text encoder with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
No EMA — targets come directly from frozen proj_in (stable, no divergence).
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import LeWMTTSv5, build_model
from dataset_v5 import build_dataloader

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

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)
    model.load_state_dict(ckpt["model"])
    print(f"Loaded pretrained from {args.pretrained} (epoch {ckpt['epoch']+1})")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    # Also freeze EMA (not used)
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder from scratch if requested
    if args.reinit:
        print("Reinitializing predictor + text encoder from scratch")
        from model_v5 import JEPAPredictor, TextEncoder
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + start_emb)")

    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, max_codec_frames=args.max_codec_frames,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(2000, 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)

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

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

    model.train()
    # Keep projections in eval mode
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

            # Text encoding (trainable)
            text_emb = model.text_encoder(text_tokens, text_mask)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()  # no EMA, direct target

            # Prepend start, shift right
            start = model.start_emb.expand(B, -1, -1)
            input_emb = torch.cat([start, z[:, :-1]], dim=1)

            # Classifier-free guidance: randomly drop audio context
            # Forces the model to use text cross-attention
            if args.audio_drop_prob > 0:
                drop_mask = torch.rand(B, device=device) < args.audio_drop_prob
                if drop_mask.any():
                    input_emb = input_emb.clone()
                    input_emb[drop_mask] = 0.0  # zero out audio for dropped samples

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Masks
            if codec_mask is not None:
                start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
                pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
                loss_mask = codec_mask
            else:
                pred_mask = None
                loss_mask = None

            # Predict
            predicted = model.predictor(input_emb, text_emb, pred_mask, text_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            if loss_mask is not None:
                valid = (~loss_mask).unsqueeze(-1)
                pred_loss = F.mse_loss(
                    predicted * valid, target_emb * valid, reduction="sum"
                ) / (valid.sum() * predicted.shape[-1] + 1e-8)
            else:
                pred_loss = F.mse_loss(predicted, target_emb)

            optimizer.zero_grad()
            pred_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} | {epoch_time:.1f}s")

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            # Save full model (including frozen projections)
            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: {ckpt_path}")

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    parser.add_argument("--audio_drop_prob", type=float, default=0.0, help="Prob of dropping audio context (CFG)")
    args = parser.parse_args()
    train(args)


Now update inference_v5.py to use duration predictor:


Write

"""
Inference for LeWM TTS v5 (codec-based JEPA with duration predictor).
Text → duration predict → expand text → JEPA predict → proj_out → EnCodec decode → waveform
"""

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

from model_v5 import LeWMTTSv5, length_regulate

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

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

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

    @torch.no_grad()
    def synthesize(self, text, duration_scale=1.0, temperature=0.0):
        """AR synthesis with duration predictor for text alignment."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Predict durations and expand text
        durations = self.model.predict_durations(text_emb)
        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
        total_frames = durations.sum().item()
        print(f"Text: {len(text)} chars → {total_frames} frames ({total_frames/75:.2f}s)")

        # Expand text to frame level
        text_emb_expanded = length_regulate(text_emb, durations)  # [1, T_total, d_model]

        # AR generation
        start_emb = self.model.start_emb.to(self.device)
        cache = self.model.init_ar_cache()
        all_embeddings = [start_emb]

        # First step
        text_frame = text_emb_expanded[:, 0:1]  # [1, 1, d_model]
        next_emb, cache = self.model.predict_next_cached(start_emb, text_frame, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

        for step in range(1, total_frames):
            # Get text embedding for this frame
            if step < text_emb_expanded.shape[1]:
                text_frame = text_emb_expanded[:, step:step+1]
            else:
                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z = self.model.proj_in(codec_emb.transpose(1, 2))
        codec_recon = self.model.proj_out(z).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → proj_in → teacher-forced predict with text alignment → proj_out → decode."""
        from encodec.utils import convert_audio
        from model_v5 import compute_uniform_durations
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        T = codec_emb.shape[2]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Compute uniform durations and expand text
        text_len = torch.tensor([text_tokens.shape[1]], device=self.device)
        audio_len = torch.tensor([T], device=self.device)
        durations = compute_uniform_durations(text_len, audio_len)
        text_emb_expanded = length_regulate(text_emb, durations)
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(1, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=self.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb_expanded)

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--duration_scale", type=float, default=1.0)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.duration_scale, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()


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

"""
Inference for LeWM TTS v5 (codec-based JEPA with duration predictor).
Text → duration predict → expand text → JEPA predict → proj_out → EnCodec decode → waveform
"""

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

from model_v5 import LeWMTTSv5, length_regulate

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

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

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

    @torch.no_grad()
    def synthesize(self, text, duration_scale=1.0, temperature=0.0):
        """AR synthesis with duration predictor for text alignment."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Predict durations and expand text
        durations = self.model.predict_durations(text_emb)
        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
        total_frames = durations.sum().item()
        print(f"Text: {len(text)} chars → {total_frames} frames ({total_frames/75:.2f}s)")

        # Expand text to frame level
        text_emb_expanded = length_regulate(text_emb, durations)  # [1, T_total, d_model]

        # AR generation
        start_emb = self.model.start_emb.to(self.device)
        cache = self.model.init_ar_cache()
        all_embeddings = [start_emb]

        # First step
        text_frame = text_emb_expanded[:, 0:1]  # [1, 1, d_model]
        next_emb, cache = self.model.predict_next_cached(start_emb, text_frame, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

        for step in range(1, total_frames):
            # Get text embedding for this frame
            if step < text_emb_expanded.shape[1]:
                text_frame = text_emb_expanded[:, step:step+1]
            else:
                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z = self.model.proj_in(codec_emb.transpose(1, 2))
        codec_recon = self.model.proj_out(z).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → proj_in → teacher-forced predict with text alignment → proj_out → decode."""
        from encodec.utils import convert_audio
        from model_v5 import compute_uniform_durations
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        T = codec_emb.shape[2]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Compute uniform durations and expand text
        text_len = torch.tensor([text_tokens.shape[1]], device=self.device)
        audio_len = torch.tensor([T], device=self.device)
        durations = compute_uniform_durations(text_len, audio_len)
        text_emb_expanded = length_regulate(text_emb, durations)
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(1, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=self.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb_expanded)

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--duration_scale", type=float, default=1.0)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.duration_scale, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()

-Inference for LeWM TTS v5 (codec-based JEPA).

-Text → JEPA predictor → upsample → EnCodec decoder → waveform

+Inference for LeWM TTS v5 (codec-based JEPA with duration predictor).

+Text → duration predict → expand text → JEPA predict → proj_out → EnCodec decode → waveform

+from model_v5 import LeWMTTSv5, length_regulate

-        self.ds_stride = config.get("ds_stride", 3)

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

-        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""

+    def synthesize(self, text, duration_scale=1.0, temperature=0.0):

+        """AR synthesis with duration predictor for text alignment."""

+        # Predict durations and expand text

+        durations = self.model.predict_durations(text_emb)

+        durations = (durations.float() * duration_scale).round().long().clamp(min=1)

+        total_frames = durations.sum().item()

+        print(f"Text: {len(text)} chars → {total_frames} frames ({total_frames/75:.2f}s)")

+        text_emb_expanded = length_regulate(text_emb, durations)  # [1, T_total, d_model]

+        cache = self.model.init_ar_cache()

-        cache = self.model.init_ar_cache()

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

+        text_frame = text_emb_expanded[:, 0:1]  # [1, 1, d_model]

+        next_emb, cache = self.model.predict_next_cached(start_emb, text_frame, 0, cache)

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

+        for step in range(1, total_frames):

+            # Get text embedding for this frame

+            if step < text_emb_expanded.shape[1]:

+                text_frame = text_emb_expanded[:, step:step+1]

+                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

-        for step in range(1, max_steps):

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

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

-            norm = next_emb.norm().item()

-                recent = prev_norms[-20:]

-                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:

-        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]

+        waveform = self.codec.decoder(codec_emb)

-        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

-        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

-        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]

+        codec_emb = self.codec.encoder(wav.unsqueeze(0))

+        z = self.model.proj_in(codec_emb.transpose(1, 2))

+        codec_recon = self.model.proj_out(z).transpose(1, 2)

-        """Encode audio → proj_in → teacher-forced predict → proj_out → decode."""

+        """Encode audio → proj_in → teacher-forced predict with text alignment → proj_out → decode."""

+        from model_v5 import compute_uniform_durations

+        # Compute uniform durations and expand text

+        text_len = torch.tensor([text_tokens.shape[1]], device=self.device)

+        audio_len = torch.tensor([T], device=self.device)

+        durations = compute_uniform_durations(text_len, audio_len)

+        if text_emb_expanded.shape[1] > T:

+            text_emb_expanded = text_emb_expanded[:, :T]

+        elif text_emb_expanded.shape[1] < T:

+            pad = torch.zeros(1, T - text_emb_expanded.shape[1], text_emb.shape[-1],

+                            device=self.device)

+            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

-        predicted = self.model.predictor(input_emb, text_emb)

+        predicted = self.model.predictor(input_emb, text_emb_expanded)

-        codec_recon = self.model.proj_out(predicted).transpose(1, 2)  # [1, 128, T]

+        codec_recon = self.model.proj_out(predicted).transpose(1, 2)

-    parser.add_argument("--max_steps", type=int, default=300)

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

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

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

waveform, sr = tts.reconstruct(args.audio)

"""
Inference for LeWM TTS v5 (codec-based JEPA).
Text → JEPA predictor → upsample → EnCodec decoder → waveform
"""

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

from model_v5 import LeWMTTSv5

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

        self.config = config
        self.ds_stride = config.get("ds_stride", 3)
        print(f"v5 model loaded from {checkpoint_path}")

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

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0):
        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # proj_in → proj_out roundtrip
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]
        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → proj_in → teacher-forced predict → proj_out → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)  # [1, 128, T]
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()


Let me verify the model builds correctly:


Bash

cd /home/ubuntu/lewm-tts && python -c "
from model_v5 import build_model, LeWMTTSv5
model, config = build_model()
print('Model built OK')
print(f'Total params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M')
print(f'Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.2f}M')
print()
# Check duration predictor
print('Duration predictor params:', sum(p.numel() for p in model.duration_predictor.parameters())/1e3, 'K')
# Test forward
import torch
B, T, T_text = 2, 100, 20
codec_emb = torch.randn(B, 128, T)
text_tokens = torch.randint(0, 256, (B, T_text))
codec_mask = torch.zeros(B, T, dtype=torch.bool)
text_mask = torch.zeros(B, T_text, dtype=torch.bool)
out = model(codec_emb, text_tokens, codec_mask, text_mask)
print('Forward OK')
for k, v in out.items():
    print(f'  {k}: {v.item():.4f}')
"

Test model builds and forward pass works


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


cd /home/ubuntu/lewm-tts && python3 -c "
from model_v5 import build_model, LeWMTTSv5
model, config = build_model()
print('Model built OK')
print(f'Total params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M')
print(f'Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.2f}M')
print()
print('Duration predictor params:', sum(p.numel() for p in model.duration_predictor.parameters())/1e3, 'K')
import torch
B, T, T_text = 2, 100, 20
codec_emb = torch.randn(B, 128, T)
text_tokens = torch.randint(0, 256, (B, T_text))
codec_mask = torch.zeros(B, T, dtype=torch.bool)
text_mask = torch.zeros(B, T_text, dtype=torch.bool)
out = model(codec_emb, text_tokens, codec_mask, text_mask)
print('Forward OK')
for k, v in out.items():
    print(f'  {k}: {v.item():.4f}')
"

Test model builds and forward pass works


LeWM TTS v5: 8.88M trainable parameters
Model built OK
Total params: 8.98M
Trainable: 8.88M

Duration predictor params: 395.009 K
Forward OK
  total_loss: 5.4811
  prediction_loss: 0.0839
  roundtrip_loss: 0.7995
  dur_loss: 3.8430
Shell cwd was reset to /home/ubuntu/pipecat


Model builds and runs. Now let's start training with reinitialized predictor:


cd /home/ubuntu/lewm-tts && python3 train_v5_pred.py \
  --pretrained output_v5/checkpoints/epoch_0120.pt \
  --output_dir output_v5_pred_dur \
  --reinit \
  --pred_weight 1.0 \
  --dur_weight 1.0 \
  --lr 5e-5 \
  --batch_size 256 \
  --epochs 200 \
  --print_every 20 \
  --save_every 10 \
  --input_noise_max 0.05 2>&1 | head -30

Start training with duration predictor (check first output)


Command running in background with ID: b67vd8o2o. Output is being written to: /tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output


sleep 15 && cat /tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output

Check training output after 15s


(Bash completed with no output)


sleep 10 && cat /tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output 2>/dev/null; ls -la /home/ubuntu/lewm-tts/output_v5_pred_dur/ 2>/dev/null

Check training status


total 20
drwxrwxr-x  4 ubuntu ubuntu 4096 Mar 24 18:44 .
drwxrwxr-x 27 ubuntu ubuntu 4096 Mar 24 18:44 ..
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 24 18:44 checkpoints
-rw-rw-r--  1 ubuntu ubuntu  265 Mar 24 18:44 config.json
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 24 18:44 logs


ToolSearch

Question: select:TaskOutput


Question: select:TaskOutput


<retrieval_status>not_ready</retrieval_status>

<task_id>b67vd8o2o</task_id>

<task_type>local_bash</task_type>

<status>running</status>

Start training with duration predictor (check first output)


Bash

cat /tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output 2>/dev/null || echo "No output yet"

Check training output directly


Check if training process is running


ubuntu   3765479  0.0  0.0   7508  3584 ?        Ss   18:44   0:00 /bin/bash -c source /home/ubuntu/.claude/shell-snapshots/snapshot-bash-1774360482847-5k7swn.sh && shopt -u extglob 2>/dev/null || true && eval 'cd /home/ubuntu/lewm-tts && python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05 2>&1 < /dev/null | head -30' && pwd -P >| /tmp/claude-4ba0-cwd
ubuntu   3765499  104  2.5 74256628 3186516 ?    Rl   18:44   1:21 python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05
ubuntu   3765547  0.9  0.5 12388704 627252 ?     Sl   18:44   0:00 python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05
ubuntu   3765548  1.1  0.5 12386732 624968 ?     Sl   18:44   0:00 python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05
ubuntu   3765549  0.8  0.6 12505064 743488 ?     Rl   18:44   0:00 python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05
ubuntu   3765550  0.7  0.5 12388564 627132 ?     Sl   18:44   0:00 python3 train_v5_pred.py --pretrained output_v5/checkpoints/epoch_0120.pt --output_dir output_v5_pred_dur --reinit --pred_weight 1.0 --dur_weight 1.0 --lr 5e-5 --batch_size 256 --epochs 200 --print_every 20 --save_every 10 --input_noise_max 0.05


Training is running but output is piped through `head -30` which buffers. Let me kill and restart without the head pipe so we can see output:


kill 3765499 2>/dev/null; sleep 2; ps aux | grep train_v5 | grep -v grep | wc -l

Kill training process


<task-notification>
<task-id>b67vd8o2o</task-id>
<tool-use-id>toolu_vrtx_011NLbQe4ZBBZBMikuDSMjGp</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output</output-file>
<status>completed</status>
<summary>Background command "Start training with duration predictor (check first output)" completed (exit code 0)</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9/tasks/b67vd8o2o.output


cd /home/ubuntu/lewm-tts && nohup python3 train_v5_pred.py \
  --pretrained output_v5/checkpoints/epoch_0120.pt \
  --output_dir output_v5_pred_dur \
  --reinit \
  --pred_weight 1.0 \
  --dur_weight 1.0 \
  --lr 5e-5 \
  --batch_size 256 \
  --epochs 200 \
  --print_every 20 \
  --save_every 10 \
  --input_noise_max 0.05 > output_v5_pred_dur/train.log 2>&1 &
echo "PID: $!"

Start training with output to log file


PID: 3766336


sleep 20 && tail -30 /home/ubuntu/lewm-tts/output_v5_pred_dur/train.log

Check training log after 20s


Device: cuda
Loaded 144 params from output_v5/checkpoints/epoch_0120.pt (skipped 36)
  Skipped: ['predictor.transformer.layers.0.multihead_attn.in_proj_weight', 'predictor.transformer.layers.0.multihead_attn.in_proj_bias', 'predictor.transformer.layers.0.multihead_attn.out_proj.weight', 'predictor.transformer.layers.0.multihead_attn.out_proj.bias', 'predictor.transformer.layers.0.norm3.weight', 'predictor.transformer.layers.0.norm3.bias', 'predictor.transformer.layers.1.multihead_attn.in_proj_weight', 'predictor.transformer.layers.1.multihead_attn.in_proj_bias', 'predictor.transformer.layers.1.multihead_attn.out_proj.weight', 'predictor.transformer.layers.1.multihead_attn.out_proj.bias']...
Reinitializing predictor + text encoder + duration predictor from scratch
Trainable params: 8.69M (predictor + text encoder + dur predictor + start_emb)
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 26 batches/epoch


Training started. The skipped params are the old cross-attention weights (which the new encoder-only predictor doesn't have) — that's expected. Let me wait for some training progress:


sleep 40 && tail -20 /home/ubuntu/lewm-tts/output_v5_pred_dur/train.log

Check training progress after ~1 minute


Still loading/first batch. The `length_regulate` loop might be slow. Let me check GPU usage:


Check GPU usage


100 %, 57843 MiB


GPU is at 100%. It's working, just the first batch is slow with the loop-based `length_regulate`. Let me wait more:


sleep 60 && tail -15 /home/ubuntu/lewm-tts/output_v5_pred_dur/train.log

Check training progress after another minute


The loop-based `length_regulate` is very slow with batch_size=256. Let me optimize it to use `torch.repeat_interleave` instead:


kill 3766336 2>/dev/null; sleep 2; ps aux | grep train_v5 | grep -v grep | wc -l

Kill slow training


kill -9 3766336 2>/dev/null; sleep 1; pkill -9 -f train_v5_pred; sleep 2; ps aux | grep train_v5 | grep -v grep | wc -l

Force kill all training processes


Exit code 1


Verify training stopped


Good, all dead. Now optimize `length_regulate` to be vectorized:


Edit

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = max(sum(durations))
    """
    B, T_text, D = text_emb.shape
    if text_mask is not None:
        durations = durations.masked_fill(text_mask, 0)
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        durs_b = durations[b]  # [T_text]
        total = durs_b.sum().item()
        if total == 0:
            continue
        # Use repeat_interleave — vectorized per sample
        exp_b = torch.repeat_interleave(text_emb[b], durs_b, dim=0)  # [total, D]
        expanded[b, :exp_b.shape[0]] = exp_b
    return expanded


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

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = max(sum(durations))
    """
    B, T_text, D = text_emb.shape
    if text_mask is not None:
        durations = durations.masked_fill(text_mask, 0)
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        durs_b = durations[b]  # [T_text]
        total = durs_b.sum().item()
        if total == 0:
            continue
        # Use repeat_interleave — vectorized per sample
        exp_b = torch.repeat_interleave(text_emb[b], durs_b, dim=0)  # [total, D]
        expanded[b, :exp_b.shape[0]] = exp_b
    return expanded

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - DurationPredictor: predicts per-character duration
  - JEPAPredictor: causal transformer with frame-level text conditioning
"""

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

# ─── Duration Predictor ─────────────────────────────────────────────────────

class DurationPredictor(nn.Module):
    """Predicts log-duration per text token. Similar to FastSpeech."""
    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
        super().__init__()
        layers = []
        for _ in range(num_layers):
            layers.extend([
                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
                nn.ReLU(),
                nn.LayerNorm(d_model),
                nn.Dropout(dropout),
            ])
        self.convs = nn.ModuleList(layers)
        self.proj = nn.Linear(d_model, 1)

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
        q = F.linear(q_input, Wq, bq)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k, v = k_new, v_new
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)
        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
        return out, k, v

    def init_cache(self, num_layers):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(ff_out))
        return self.output_proj(x), cache

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

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

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        return self.proj(x)

# ─── Full v5 Model ──────────────────────────────────────────────────────────

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend + duration predictor.
    Text is expanded to frame level and added to audio input — can't be ignored.

    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
              Duration predictor trained with uniform alignment as supervision.
    Inference: text → duration predict → expand text → AR predict → proj_out → decode
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Duration predictor
        self.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

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

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
        self.dur_weight = config.get("dur_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Duration prediction + ground truth durations (uniform alignment)
        log_dur_pred = self.duration_predictor(text_emb, text_mask)

        # Compute actual text/audio lengths from masks
        if text_mask is not None:
            text_lengths = (~text_mask).sum(dim=1)  # [B]
        else:
            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)
        if codec_mask is not None:
            audio_lengths = (~codec_mask).sum(dim=1)  # [B]
        else:
            audio_lengths = torch.full((B,), T, device=codec_emb.device)

        # Ground truth durations (uniform distribution)
        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

        # Duration loss (MSE on log-durations)
        if text_mask is not None:
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)
        else:
            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)

        # Expand text to frame level using GT durations (training uses GT)
        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
        # Trim/pad to match audio length T
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=text_emb.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

        # JEPA prediction with frame-level text
        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = (self.pred_weight * prediction_loss
                     + self.roundtrip_weight * roundtrip_loss
                     + self.dur_weight * dur_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
            "dur_loss": dur_loss,
        }

    def predict_durations(self, text_emb, text_mask=None):
        """Predict durations from text embeddings (for inference)."""
        log_dur = self.duration_predictor(text_emb, text_mask)
        durations = torch.exp(log_dur).round().long().clamp(min=1)
        if text_mask is not None:
            durations = durations.masked_fill(text_mask, 0)
        return durations

    def init_ar_cache(self):
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model]
        text_emb_frame: [B, 1, d_model] — expanded text for this frame
        """
        return self.predictor.inference_step(new_emb, text_emb_frame, step_idx, cache)

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

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256, "nhead": 4, "codec_dim": 128,
            "text_vocab_size": 256, "text_encoder_layers": 4,
            "predictor_layers": 6, "dropout": 0.1,
            "pred_weight": 10.0, "roundtrip_weight": 1.0,
            "ema_decay": 0.998, "input_noise": 0.0,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

-    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample

+    Returns: expanded [B, T_audio, d_model], where T_audio = max(sum(durations))

+        durations = durations.masked_fill(text_mask, 0)

-            if text_mask is not None and text_mask[b, t]:

-            dur = durations[b, t].item()

-            if dur > 0 and pos < max_audio_len:

-                end = min(pos + dur, max_audio_len)

-                expanded[b, pos:end] = text_emb[b, t]

+        durs_b = durations[b]  # [T_text]

+        # Use repeat_interleave — vectorized per sample

+        exp_b = torch.repeat_interleave(text_emb[b], durs_b, dim=0)  # [total, D]

+        expanded[b, :exp_b.shape[0]] = exp_b


Also optimize `compute_uniform_durations` to be vectorized:


def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    # Vectorized per sample (text lengths vary)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations