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

            # 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)
    args = parser.parse_args()
    train(args)

+    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")


Bash

cd /home/ubuntu/lewm-tts && rm -rf output_v5_pred && nohup python3 -u train_v5_pred.py \
    --pretrained output_v5/checkpoints/epoch_0120.pt \
    --data_dir processed_data_codec \
    --output_dir output_v5_pred \
    --epochs 200 \
    --batch_size 256 \
    --lr 1e-4 \
    --input_noise_max 0.05 \
    --max_codec_frames 900 \
    --num_workers 4 \
    --print_every 20 \
    --save_every 10 \
    --reinit \
    > train_v5_pred.log 2>&1 &
echo "PID: $!"

Launch predictor training with fresh initialization


PID: 3735571


sleep 90 && tail -15 /home/ubuntu/lewm-tts/train_v5_pred.log

Check fresh predictor training


Device: cuda
Loaded pretrained from output_v5/checkpoints/epoch_0120.pt (epoch 120)
Reinitializing predictor + text encoder from scratch
Trainable params: 9.81M (predictor + text encoder + start_emb)
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 26 batches/epoch
  [1/200] step 20 | pred=11.1494 | grad=3.86 | 29.4s
Epoch 1/200 | pred=11.2827 | 37.6s
  New best (pred=11.2827)
  [2/200] step 40 | pred=10.6935 | grad=3.39 | 20.0s
Epoch 2/200 | pred=10.8558 | 36.8s
  New best (pred=10.8558)
  [3/200] step 60 | pred=10.4425 | grad=3.70 | 11.7s


Starting from 11.28 (random init) and dropping. Much better gradient flow now (grad=3.86 vs 0.27 before). The predictor is actually learning. Let me check after a few more epochs:


sleep 300 && tail -10 /home/ubuntu/lewm-tts/train_v5_pred.log

Check after 5 minutes


Epoch 8/200 | pred=6.4919 | 36.8s
  New best (pred=6.4919)
  [9/200] step 220 | pred=5.2533 | grad=6.19 | 17.3s
Epoch 9/200 | pred=5.1514 | 37.0s
  New best (pred=5.1514)
  [10/200] step 240 | pred=4.1528 | grad=5.81 | 8.8s
  [10/200] step 260 | pred=3.1641 | grad=5.14 | 36.9s
Epoch 10/200 | pred=3.7784 | 37.0s
  Saved: output_v5_pred/checkpoints/epoch_0010.pt
  New best (pred=3.7784)


Now **that's** a learning curve! 11.28 → 3.78 in 10 epochs. Compare to the stuck predictor before (0.405 → 0.397 in 10 epochs).

The loss should keep dropping. At ~37s/epoch, I'll generate samples around epoch 50 when it's had time to converge. That'll be about 25 minutes from now.


status


Check current training


  New best (pred=0.7167)
  [16/200] step 400 | pred=0.6433 | grad=0.29 | 14.5s
Epoch 16/200 | pred=0.6381 | 36.9s
  New best (pred=0.6381)
  [17/200] step 420 | pred=0.5992 | grad=0.69 | 6.0s


Epoch 17. Loss: **11.28 → 0.60** in 17 epochs. Dropping fast and still converging. Compare to the stuck run before (0.40 → 0.40).

About 20 more minutes until epoch 50 for samples.


create samples on this epoch


ls -t /home/ubuntu/lewm-tts/output_v5_pred/checkpoints/ | head -3

Check available checkpoints


best.pt
epoch_0010.pt


Only epoch 10 saved. Let me use `best.pt` (latest best):


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()
print(f'Loaded epoch {ckpt[\"epoch\"]+1}, pred_loss={ckpt[\"loss\"]:.4f}')

codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5_pred', exist_ok=True)

def save_wav(w, path, sr=24000):
    mx = np.abs(w).max()
    if mx > 0: w = w / mx * 0.95
    torchaudio.save(path, torch.from_numpy(w).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(w)/sr:.2f}s)')

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

entry = manifest[0]
text, audio_path = entry['text'], entry['audio_path']
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)

with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))

    # GT + codec baseline
    save_wav(wav_24k.squeeze().numpy(), 'samples_v5_pred/gt.wav')
    save_wav(codec.decoder(codec_emb).squeeze().cpu().numpy(), 'samples_v5_pred/codec_only.wav')

    # Roundtrip (should still be good — projections frozen)
    z = model.proj_in(codec_emb.transpose(1, 2))
    codec_recon = model.proj_out(z).transpose(1, 2)
    print(f'Roundtrip L1: {(codec_emb - codec_recon).abs().mean():.4f}')
    save_wav(codec.decoder(codec_recon).squeeze().cpu().numpy(), 'samples_v5_pred/roundtrip.wav')

    # Teacher-forced
    text_emb = model.text_encoder(text_to_tokens(text))
    start = model.start_emb.expand(1, -1, -1)
    input_emb = torch.cat([start, z[:, :-1]], dim=1)
    predicted = model.predictor(input_emb, text_emb)
    codec_tf = model.proj_out(predicted).transpose(1, 2)
    print(f'TF L1: {(codec_emb - codec_tf).abs().mean():.4f}')
    save_wav(codec.decoder(codec_tf).squeeze().cpu().numpy(), 'samples_v5_pred/tf.wav')

    # AR synth temp=0
    all_embs = [model.start_emb.to(device)]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(all_embs[0], text_emb, 0, cache)
    all_embs.append(next_emb)
    for step in range(1, 750):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        all_embs.append(next_emb)
        if step > 50:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.proj_out(pred_embs).transpose(1, 2)
    save_wav(codec.decoder(codec_synth).squeeze().cpu().numpy(), 'samples_v5_pred/synth_t00.wav')
    print(f'Synth t=0: {step+1} steps, {pred_embs.shape[1]} embs')

    # AR synth temp=0.3
    all_embs = [model.start_emb.to(device)]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(all_embs[0], text_emb, 0, cache)
    next_emb = next_emb + torch.randn_like(next_emb) * 0.3
    all_embs.append(next_emb)
    for step in range(1, 750):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        next_emb = next_emb + torch.randn_like(next_emb) * 0.3
        all_embs.append(next_emb)
        if step > 50:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.proj_out(pred_embs).transpose(1, 2)
    save_wav(codec.decoder(codec_synth).squeeze().cpu().numpy(), 'samples_v5_pred/synth_t03.wav')
    print(f'Synth t=0.3: {step+1} steps')

print('\\nAll in samples_v5_pred/')
" 2>&1

Generate samples from latest best checkpoint


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loaded epoch 17, pred_loss=0.5985
Saved: samples_v5_pred/gt.wav (9.52s)
Saved: samples_v5_pred/codec_only.wav (9.52s)
Roundtrip L1: 0.2699
Saved: samples_v5_pred/roundtrip.wav (9.52s)
TF L1: 0.7180
Saved: samples_v5_pred/tf.wav (9.52s)
Saved: samples_v5_pred/synth_t00.wav (1.47s)
Synth t=0: 109 steps, 110 embs
Saved: samples_v5_pred/synth_t03.wav (6.64s)
Synth t=0.3: 497 steps

All in samples_v5_pred/
Shell cwd was reset to /home/ubuntu/pipecat


Samples in `samples_v5_pred/`:

| File | Description |
|------|-------------|
| `gt.wav` | Ground truth (9.52s) |
| `codec_only.wav` | Pure EnCodec ceiling (9.52s) |
| `roundtrip.wav` | Projection roundtrip — still good, L1=0.27 (9.52s) |
| `tf.wav` | Teacher-forced, TF L1=0.72 (9.52s) |
| `synth_t00.wav` | AR temp=0 (1.47s — still collapses) |
| `synth_t03.wav` | AR temp=0.3 (6.64s — longer than before!) |

Pred loss is at 0.60 — still early (was 11.28 at start). The TF L1=0.72 is higher than the previous predictor's 0.61 because this one just started training. But it's learning fast.

Give these a listen — especially `tf.wav` compared to the robotic one from before. The synth at temp=0.3 is producing 6.6s now (was 0.71s before).


only noise . it is trying to speak i think .. when person speaking at that time noise cmng instead of voice


The user says only noise is coming, but they can tell it's trying to speak - noise comes when the person should be speaking. This means the predictor is capturing the temporal pattern (when speech vs silence) but the actual embedding values are wrong.

At pred_loss=0.60 with epoch 17, the predictor is still very early in training. The loss started at 11.28 and is at 0.60. It needs to go much lower - probably below 0.1 to get intelligible speech.

But there's a deeper issue: at 75Hz, the sequences are very long (700+ frames). The 6-layer transformer predictor with d_model=256 might not have enough capacity, or the causal attention over 700 tokens is very expensive and the model can't learn long-range dependencies well.

Actually, I'm comparing this to v4 where the same architecture reached 0.073 loss after 120 epochs at 23Hz. The difference is that 75Hz gives me 3x more frames per utterance, and with quadratic attention complexity, the gradient signal gets spread thinner across more tokens. The loss is dropping quickly though - it went from 11.28 to 0.60 in just 17 epochs, so maybe I should just let it continue training and see if it reaches comparable loss values by epoch 100-200.

But the user is getting impatient, and realistically I need the loss to drop below ~0.1 for intelligible speech synthesis. At the current rate that could take another 100+ epochs. I should probably tell them we're only 8.5% through training and the trajectory looks promising, but I'm also noticing something odd - the TF L1 loss is 0.72, which is actually higher than the old predictor's 0.61, so something might be off with how the embeddings are scaling in this 256d space.


That's expected at epoch 17 — the predictor literally just started (reinitialized from scratch 17 epochs ago). The loss went 11.28 → 0.60 but needs to get below ~0.1 for intelligible speech (that's where v4's predictor was).

It's 8.5% through training. Let me wait until epoch 50+ and check again. Should be about 20 minutes.


sleep 1200 && tail -5 /home/ubuntu/lewm-tts/train_v5_pred.log

Wait 20 min then check


check now


Error: Exit code 137
[Request interrupted by user for tool use]


[Request interrupted by user for tool use]


Check training progress


  [43/200] step 1100 | pred=0.4380 | grad=0.21 | 11.5s
Epoch 43/200 | pred=0.4372 | 37.0s
  New best (pred=0.4372)
  [44/200] step 1120 | pred=0.4457 | grad=0.30 | 3.2s
  [44/200] step 1140 | pred=0.4393 | grad=0.15 | 31.1s


Epoch 44, pred_loss=0.44. It's plateauing — was 0.60 at epoch 17, now 0.44 at epoch 44. Barely moving again. Same problem as before — the predictor is hitting a floor around 0.4.

This tells me the 6-layer predictor at 75Hz can't get below ~0.4 MSE. The architecture isn't powerful enough for 75Hz prediction. In v4 at 23Hz it reached 0.07 — that's 6x lower.

The honest assessment: **predicting continuous codec embeddings at 75Hz with this architecture won't get to good quality.** We either need:
1. Much larger predictor (more layers/dims) — but memory limited
2. Lower frame rate (downsample, but that caused noise)
3. **Switch to discrete token prediction** — the proven approach

Want me to try discrete token prediction with SNAC (12Hz, level 0)? That's what actually works in production TTS systems. Or should we keep pushing on this approach?


can u genarte some samples


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()
print(f'Loaded epoch {ckpt[\"epoch\"]+1}, pred_loss={ckpt[\"loss\"]:.4f}')

codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5_pred_ep44', exist_ok=True)

def save_wav(w, path, sr=24000):
    mx = np.abs(w).max()
    if mx > 0: w = w / mx * 0.95
    torchaudio.save(path, torch.from_numpy(w).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(w)/sr:.2f}s)')

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

entry = manifest[0]
text, audio_path = entry['text'], entry['audio_path']
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)

with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    z = model.proj_in(codec_emb.transpose(1, 2))
    text_emb = model.text_encoder(text_to_tokens(text))

    # Roundtrip
    codec_recon = model.proj_out(z).transpose(1, 2)
    print(f'Roundtrip L1: {(codec_emb - codec_recon).abs().mean():.4f}')
    save_wav(codec.decoder(codec_recon).squeeze().cpu().numpy(), 'samples_v5_pred_ep44/roundtrip.wav')

    # TF
    start = model.start_emb.expand(1, -1, -1)
    input_emb = torch.cat([start, z[:, :-1]], dim=1)
    predicted = model.predictor(input_emb, text_emb)
    codec_tf = model.proj_out(predicted).transpose(1, 2)
    print(f'TF L1: {(codec_emb - codec_tf).abs().mean():.4f}')
    save_wav(codec.decoder(codec_tf).squeeze().cpu().numpy(), 'samples_v5_pred_ep44/tf.wav')

    # AR synth temp=0
    all_embs = [model.start_emb.to(device)]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(all_embs[0], text_emb, 0, cache)
    all_embs.append(next_emb)
    for step in range(1, 750):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        all_embs.append(next_emb)
        if step > 50:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.proj_out(pred_embs).transpose(1, 2)
    save_wav(codec.decoder(codec_synth).squeeze().cpu().numpy(), 'samples_v5_pred_ep44/synth_t00.wav')
    print(f'Synth t=0: {step+1} steps')

    # AR synth temp=0.3
    all_embs = [model.start_emb.to(device)]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(all_embs[0], text_emb, 0, cache)
    next_emb = next_emb + torch.randn_like(next_emb) * 0.3
    all_embs.append(next_emb)
    for step in range(1, 750):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        next_emb = next_emb + torch.randn_like(next_emb) * 0.3
        all_embs.append(next_emb)
        if step > 50:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.proj_out(pred_embs).transpose(1, 2)
    save_wav(codec.decoder(codec_synth).squeeze().cpu().numpy(), 'samples_v5_pred_ep44/synth_t03.wav')
    print(f'Synth t=0.3: {step+1} steps')

print('\\nAll in samples_v5_pred_ep44/')
" 2>&1

Generate samples from epoch 44


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loaded epoch 45, pred_loss=0.4341
Roundtrip L1: 0.2699
Saved: samples_v5_pred_ep44/roundtrip.wav (9.52s)
TF L1: 0.6226
Saved: samples_v5_pred_ep44/tf.wav (9.52s)
Saved: samples_v5_pred_ep44/synth_t00.wav (1.51s)
Synth t=0: 112 steps
Saved: samples_v5_pred_ep44/synth_t03.wav (1.65s)
Synth t=0.3: 123 steps

All in samples_v5_pred_ep44/
Shell cwd was reset to /home/ubuntu/pipecat


Samples in `samples_v5_pred_ep44/`:

| File | Description |
|------|-------------|
| `roundtrip.wav` | Projection roundtrip (9.52s) — still good |
| `tf.wav` | Teacher-forced, TF L1=0.62 (improved from 0.72) |
| `synth_t00.wav` | AR temp=0 (1.51s) |
| `synth_t03.wav` | AR temp=0.3 (1.65s) |

TF improved (0.72→0.62) but synth still collapses to ~1.5s. Have a listen to `tf.wav` — that shows the predictor quality with real inputs.


tf better than previous one


The user says TF is better than the previous one. So the predictor IS improving, just slowly. The loss went from 0.60 to 0.44 and TF L1 from 0.72 to 0.62. The quality is improving.

The question is: will it get good enough? The loss is plateauing around 0.43. Let me check if it's still dropping or truly stuck.

Actually, let me look at the trend more carefully. The loss went:
- Epoch 1: 11.28
- Epoch 17: 0.60
- Epoch 44: 0.44

It's still dropping but very slowly.

If the trend continues, I might see it reach around 0.3 by epoch 200, which could give intelligible TF output but probably won't be clean enough for AR synthesis. The user seems cautiously optimistic, so I should let training continue and check back at epoch 100. The synth collapse from AR prediction divergence at 75Hz is a separate problem—if the TF keeps improving, that's at least progress on one front.


Good sign! TF is improving as the predictor trains (L1: 0.72 → 0.62). The loss is still dropping, just slowly now. Let it keep training — by epoch 100-200 the TF should be cleaner.

The synth collapse is a separate issue (error accumulation over 750 steps at 75Hz), but if TF gets good enough, we can address that later with techniques like classifier-free guidance or beam search.

I'll check back around epoch 100 with new samples. That's about 35 minutes from now. Want me to wait and generate then, or anything else in the meantime?


but synth is not good y


The user is asking why synthesis isn't working. Let me think about this:

1. **TF (teacher-forced)**: Uses REAL encoded embeddings as input, predicts next step. Each prediction has the correct previous embedding as context. Quality is improving.

2. **Synthesis (AR)**: Uses its OWN predictions as input. Each step feeds the previous prediction back. Errors compound over steps.

At 75Hz, a 10-second utterance needs ~750 AR steps. Even if each step has 1% error, after 750 steps the errors compound massively. At 23Hz (v4), it was only ~230 steps.

The core issues with AR synth:

- Small errors accumulate across hundreds of steps, and the predictor tends to collapse toward a fixed point because it's seeing its own degraded outputs and can't recover
- There's a fundamental mismatch: training uses real embeddings as input, but inference uses predicted ones—classic exposure bias
- The noise we're adding during training (0.05 magnitude) helps but isn't enough to bridge that gap at 750 steps

I'm thinking scheduled sampling could help—using the model's own predictions during training sometimes instead of always relying on ground truth.

But stepping back, 75Hz is just too aggressive for autoregressive generation. The models that actually work in this space either operate at much lower frame rates (like 12Hz with codec tokens) or abandon the autoregressive approach entirely in favor of non-autoregressive or diffusion-based methods that generate everything in parallel.

I should be straightforward with the user about these fundamental constraints and walk through what's actually feasible.


The synth problem is **error accumulation at 75Hz**.

Here's why:
- **TF (teacher-forced)**: Each step gets the **real** previous embedding as input → small errors stay small
- **Synth (AR)**: Each step gets its **own prediction** as input → errors compound step after step

At 75Hz, a 10-second utterance needs **~750 AR steps**. Even 1% error per step becomes catastrophic after 750 steps — the signal drifts until it collapses.

For comparison, v4 at 23Hz needed only ~230 steps. And successful TTS models like VALL-E use ~75 steps (EnCodec level 0 at 7.5Hz) or SNAC at 12Hz (~120 steps).

**750 steps of AR is just too many for a small model.**

Options to fix:
1. **SNAC at 12Hz** → only ~120 AR steps for 10s (discrete tokens, proven to work)
2. **Non-autoregressive** → predict all frames at once (MaskGIT/SoundStorm style)
3. **Scheduled sampling** → train with the model's own predictions instead of always ground truth (helps but doesn't solve the fundamental 750-step problem)

The TF improving proves the predictor can learn the right patterns. The issue is purely AR at 75Hz being too many steps. Want me to try SNAC at 12Hz? That's the most proven path.


cant we use tf at infernce ?


The user is asking if we can use teacher-forced mode at inference. In TF, we feed the real encoded embeddings as input to the predictor. But at inference, we don't have the real audio — that's what we're trying to generate!

However, there IS a way to use TF-like inference: **voice cloning / voice prompting**. If we have a short reference audio clip (say 2-3 seconds), we can:
1. Encode the reference audio through EnCodec → get real codec embeddings
2. Project through proj_in → get real latent embeddings
3. Feed these as the prefix/prompt to the predictor
4. Let the predictor continue from there with the new text

This approach is essentially "prompted" generation — the model anchors on real audio and only predicts a continuation, which limits error accumulation since the initial context is grounded and the AR steps are fewer. It's how systems like VALL-E and XTTS work, using a reference audio as a prompt to guide generation. The user might be wondering if we can use real audio during inference for new speech generation, but that only works for reconstruction or modification, not synthesis from text — so voice prompting is the practical solution here.


Smart thinking! You can't use TF directly at inference (you don't have the audio you're trying to generate). But you CAN use **voice prompting** — which is TF-like:

1. Take a **short reference audio** (2-3 seconds of the target voice)
2. Encode it through EnCodec → real codec embeddings
3. Feed those as the **prefix** to the predictor (TF mode for the prompt)
4. AR-generate only the **continuation** with new text

This way:
- The model starts from **real embeddings** (not random start_emb)
- It only needs to AR-predict a few seconds beyond the prompt
- Error accumulation is limited since the context is grounded

Let me test this right now:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()

codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5_prompted', exist_ok=True)

def save_wav(w, path, sr=24000):
    mx = np.abs(w).max()
    if mx > 0: w = w / mx * 0.95
    torchaudio.save(path, torch.from_numpy(w).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(w)/sr:.2f}s)')

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

entry = manifest[0]
text, audio_path = entry['text'], entry['audio_path']
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)

with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))  # [1, 128, T]
    z_full = model.proj_in(codec_emb.transpose(1, 2))  # [1, T, 256]
    text_emb = model.text_encoder(text_to_tokens(text))
    T_full = z_full.shape[1]

    # Use first N seconds as prompt, AR-generate the rest
    for prompt_secs in [2, 3, 5]:
        prompt_frames = int(prompt_secs * 75)  # 75Hz
        prompt_z = z_full[:, :prompt_frames]  # [1, prompt_frames, 256]

        # Start with real embeddings, then AR from there
        # First: TF over the prompt (build up KV cache with real context)
        all_embs = []
        cache = model.init_ar_cache()

        # Feed prompt through predictor one-by-one to build cache
        start = model.start_emb.to(device)
        out, cache = model.predict_next_cached(start, text_emb, 0, cache)
        all_embs.append(out)  # predicted first frame

        for i in range(prompt_frames - 1):
            # Feed REAL embedding (TF), predict next
            out, cache = model.predict_next_cached(prompt_z[:, i:i+1], text_emb, i + 1, cache)
            all_embs.append(out)

        # Now AR-generate continuation using last real embedding
        next_emb = prompt_z[:, -1:]  # last real embedding
        ar_steps = 0
        for step in range(prompt_frames, prompt_frames + 375):  # up to 5 more seconds
            out, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
            all_embs.append(out)
            next_emb = out  # AR: use own prediction
            ar_steps += 1
            if ar_steps > 50:
                norms = [e.norm().item() for e in all_embs[-20:]]
                if np.std(norms) < 0.01 * np.mean(norms): break

        # Replace prompt predictions with real embeddings (TF quality)
        real_part = prompt_z  # [1, prompt_frames, 256]
        ar_part = torch.cat(all_embs[prompt_frames:], dim=1)  # [1, ar_steps, 256]
        combined = torch.cat([real_part, ar_part], dim=1)

        codec_out = model.proj_out(combined).transpose(1, 2)
        wav_out = codec.decoder(codec_out).squeeze().cpu().numpy()

        total_secs = len(wav_out) / 24000
        save_wav(wav_out, f'samples_v5_prompted/prompted_{prompt_secs}s.wav')
        print(f'  Prompt: {prompt_secs}s real + {ar_steps} AR steps = {total_secs:.2f}s total')

    # Also save the ground truth and the prompt audio for reference
    save_wav(wav_24k.squeeze().numpy(), 'samples_v5_prompted/gt.wav')

    # Save just the 3s prompt for reference
    prompt_wav = wav_24k[:, :3*24000]
    save_wav(prompt_wav.squeeze().numpy(), 'samples_v5_prompted/prompt_3s.wav')

print('\\nAll in samples_v5_prompted/')
" 2>&1

Generate voice-prompted samples with real audio prefix


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Saved: samples_v5_prompted/prompted_2s.wav (2.68s)
  Prompt: 2s real + 51 AR steps = 2.68s total
Saved: samples_v5_prompted/prompted_3s.wav (3.68s)
  Prompt: 3s real + 51 AR steps = 3.68s total
Saved: samples_v5_prompted/prompted_5s.wav (5.68s)
  Prompt: 5s real + 51 AR steps = 5.68s total
Saved: samples_v5_prompted/gt.wav (9.52s)
Saved: samples_v5_prompted/prompt_3s.wav (3.00s)

All in samples_v5_prompted/
Shell cwd was reset to /home/ubuntu/pipecat


Samples in `samples_v5_prompted/`:

| File | Description |
|------|-------------|
| `gt.wav` | Ground truth (9.52s) |
| `prompt_3s.wav` | Just the 3s prompt audio |
| `prompted_2s.wav` | 2s real prompt + AR continuation (2.68s total) |
| `prompted_3s.wav` | 3s real prompt + AR continuation (3.68s total) |
| `prompted_5s.wav` | 5s real prompt + AR continuation (5.68s total) |

The AR continuation only ran ~51 steps before collapsing (0.68s of generated audio). Still early in training. But listen to these — the prompt part should sound clean (it's real embeddings through roundtrip), and the transition to AR should be smoother than pure synthesis since the context is grounded.

Compare `prompted_3s.wav` to the pure `synth` samples — should be much better for the prompt portion at least.


these are fucking good


The user says the prompted samples are "fucking good"! The voice prompting approach works!

Now the issue is that the AR continuation only runs ~51 steps (0.68s) before collapsing. We need to:
1. Keep training the predictor so it can sustain longer AR generations
2. The prompted approach is the way forward

This validates the whole pipeline:
- EnCodec encoder → proj_in → proj_out → EnCodec decoder: WORKS (roundtrip good)
- TF prediction: WORKS (improving)
- Voice-prompted generation: WORKS ("fucking good")
- Pure AR synthesis: Still collapses (needs more training + possibly architectural fixes)

- As the predictor trains further and loss decreases, the AR continuation should stabilize and extend beyond the current 51-step limit. The prompted inference mode is what we'll use in production while the pure AR synthesis improves with more training.


Let's go! The voice prompting approach works. The pipeline is validated:
- **Roundtrip**: good
- **TF**: improving
- **Voice-prompted**: fucking good

The only issue is the AR continuation collapses after ~51 steps (0.68s). As the predictor keeps training (epoch 45/200), the continuation will get longer and more stable.

Let me let the training run and generate new prompted samples at later epochs to see the AR part extend. Training should finish in about 1.5 hours.


Edit

## v5 Architecture (codec-based)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)
- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)
- EMA on downsampler, learnable start_emb
- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)
- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100
- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## v5 Architecture (codec-based, simplified)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- Simple linear projections: proj_in(128→256→256), proj_out(256→256→128) — NO downsampler
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder) at 75Hz
- Projections trained first (output_v5/), then frozen for predictor training (output_v5_pred/)
- Predictor reinitialized from scratch with frozen proj_in targets (no EMA — EMA caused divergence)
- Files: model_v5.py, train_v5.py, train_v5_pred.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## Key Findings
- Down/upsampler CNN approach: FAILED (noisy roundtrip)
- EMA on small proj_in: FAILED (pred_loss diverged 0.12→0.38)
- Voice prompting works great: encode real audio as prefix, AR-continue from there
- Pure AR at 75Hz collapses after ~50 steps (too many steps for small model)
- Roundtrip (proj_in→proj_out→decode): clean audio, L1=0.27
- TF prediction: improving, L1=0.62 at epoch 45


The file /home/ubuntu/.claude/projects/-home-ubuntu-pipecat/memory/lewm-tts.md has been updated successfully.

## v5 Architecture (codec-based, simplified)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- Simple linear projections: proj_in(128→256→256), proj_out(256→256→128) — NO downsampler
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder) at 75Hz
- Projections trained first (output_v5/), then frozen for predictor training (output_v5_pred/)
- Predictor reinitialized from scratch with frozen proj_in targets (no EMA — EMA caused divergence)
- Files: model_v5.py, train_v5.py, train_v5_pred.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## Key Findings
- Down/upsampler CNN approach: FAILED (noisy roundtrip)
- EMA on small proj_in: FAILED (pred_loss diverged 0.12→0.38)
- Voice prompting works great: encode real audio as prefix, AR-continue from there
- Pure AR at 75Hz collapses after ~50 steps (too many steps for small model)
- Roundtrip (proj_in→proj_out→decode): clean audio, L1=0.27
- TF prediction: improving, L1=0.62 at epoch 45

# LeWM TTS - JEPA-based Hindi TTS

## Key Issues Found (v3/v3b/v3c)
1. Missing EMA target encoder (core JEPA component)
2. Prediction loss was <1% of total loss — predictor learned identity mapping
3. KL divergence exploded to 2.45 (kl_weight=0.01 too low)
4. v3c's downsample_factor=2 made consecutive embeddings too similar
5. Random start embedding at inference (never trained)
6. No input noise to bridge train/inference gap

## v4 Fixes Applied
- **EMA target encoder**: deepcopy of AudioEncoder, updated with momentum 0.998
- **Rebalanced losses**: pred_weight=10.0 (now ~36% of total), kl_weight=0.05
- **Free-bits KL**: threshold=1.0, prevents both collapse and explosion
- **Learnable start_emb**: prepended to training sequence, receives gradients
- **Input noise**: annealed 0→0.2 over first 30% of training
- **Downsample factor**: back to 4 (more temporal separation)

## v4 Results
- MelDecoder was the bottleneck (L1=0.57 even with perfect embeddings)
- JEPA predictor worked well: cosine_sim=0.97, 6.2x better than copy baseline
- Recon had noise overlay, synthesis was pure noise

## v5 Architecture (codec-based)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)
- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)
- EMA on downsampler, learnable start_emb
- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)
- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100
- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## Dataset
7009 segments, 17.88h single-speaker Hindi, processed_data_100mel/ and processed_data_codec/

+## v5 Architecture (codec-based, simplified)

-- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)

-- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)

-- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)

-- EMA on downsampler, learnable start_emb

-- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)

-- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100

-- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py

+- Simple linear projections: proj_in(128→256→256), proj_out(256→256→128) — NO downsampler

+- TextEncoder(4L,256d) + JEPAPredictor(6L decoder) at 75Hz

+- Projections trained first (output_v5/), then frozen for predictor training (output_v5_pred/)

+- Predictor reinitialized from scratch with frozen proj_in targets (no EMA — EMA caused divergence)

+- Files: model_v5.py, train_v5.py, train_v5_pred.py, dataset_v5.py, inference_v5.py, preprocess_codec.py

+- Down/upsampler CNN approach: FAILED (noisy roundtrip)

+- EMA on small proj_in: FAILED (pred_loss diverged 0.12→0.38)

+- Voice prompting works great: encode real audio as prefix, AR-continue from there

+- Pure AR at 75Hz collapses after ~50 steps (too many steps for small model)

+- Roundtrip (proj_in→proj_out→decode): clean audio, L1=0.27

+- TF prediction: improving, L1=0.62 at epoch 45


on what checkpoint u ran prompted?


Training is still running. Want me to check back at epoch 100 with new prompted samples to see if the AR continuation extends longer?


The prompted samples were from **epoch 45** (`output_v5_pred/checkpoints/best.pt`, pred_loss=0.4341).

That's the predictor trained from scratch with frozen projections from v5 epoch 120. Still very early — only 22% through training.


can u test other text not this one


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
import json
manifest = json.load(open('processed_data_codec/manifest.json'))
# Show a few different samples with their text and duration
for i in [0, 50, 100, 200, 500]:
    e = manifest[i]
    print(f'Sample {i}: {e[\"duration\"]:.1f}s | {e[\"text\"][:80]}...')
" 2>&1

Find different text samples


Interrupted by user


[Request interrupted by user]


some neew text which is not is dataset


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()

codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5_newtext', exist_ok=True)

def save_wav(w, path, sr=24000):
    mx = np.abs(w).max()
    if mx > 0: w = w / mx * 0.95
    torchaudio.save(path, torch.from_numpy(w).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(w)/sr:.2f}s)')

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

# Use sample 0's audio as voice prompt (3s), but with NEW text not in dataset
audio_path = manifest[0]['audio_path']
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)

new_texts = [
    'आज मौसम बहुत अच्छा है और हम सब खुश हैं',
    'भारत एक महान देश है जहां अनेक भाषाएं बोली जाती हैं',
    'नमस्ते दोस्तों, आज हम बात करेंगे नई तकनीक के बारे में',
]

with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    z_full = model.proj_in(codec_emb.transpose(1, 2))

    prompt_frames = int(3 * 75)  # 3 seconds
    prompt_z = z_full[:, :prompt_frames]

    for idx, new_text in enumerate(new_texts):
        text_emb = model.text_encoder(text_to_tokens(new_text))

        # Build cache with real prompt
        all_embs = []
        cache = model.init_ar_cache()
        start = model.start_emb.to(device)
        out, cache = model.predict_next_cached(start, text_emb, 0, cache)
        all_embs.append(out)

        for i in range(prompt_frames - 1):
            out, cache = model.predict_next_cached(prompt_z[:, i:i+1], text_emb, i + 1, cache)
            all_embs.append(out)

        # AR continuation
        next_emb = prompt_z[:, -1:]
        for step in range(prompt_frames, prompt_frames + 375):
            out, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
            all_embs.append(out)
            next_emb = out
            if step > prompt_frames + 50:
                norms = [e.norm().item() for e in all_embs[-20:]]
                if np.std(norms) < 0.01 * np.mean(norms): break

        real_part = prompt_z
        ar_part = torch.cat(all_embs[prompt_frames:], dim=1)
        combined = torch.cat([real_part, ar_part], dim=1)

        codec_out = model.proj_out(combined).transpose(1, 2)
        wav_out = codec.decoder(codec_out).squeeze().cpu().numpy()

        save_wav(wav_out, f'samples_v5_newtext/new_{idx}.wav')
        print(f'  Text: {new_text}')
        print(f'  AR steps: {step - prompt_frames + 1}')
        print()

print('All in samples_v5_newtext/')
" 2>&1

Generate prompted samples with new unseen text


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Saved: samples_v5_newtext/new_0.wav (3.69s)
  Text: आज मौसम बहुत अच्छा है और हम सब खुश हैं
  AR steps: 52

Saved: samples_v5_newtext/new_1.wav (3.69s)
  Text: भारत एक महान देश है जहां अनेक भाषाएं बोली जाती हैं
  AR steps: 52

Saved: samples_v5_newtext/new_2.wav (3.69s)
  Text: नमस्ते दोस्तों, आज हम बात करेंगे नई तकनीक के बारे में
  AR steps: 52

All in samples_v5_newtext/
Shell cwd was reset to /home/ubuntu/pipecat


Samples in `samples_v5_newtext/` — 3s voice prompt + AR continuation with **new Hindi text not in the dataset**:

| File | Text |
|------|------|
| `new_0.wav` | आज मौसम बहुत अच्छा है और हम सब खुश हैं |
| `new_1.wav` | भारत एक महान देश है जहां अनेक भाषाएं बोली जाती हैं |
| `new_2.wav` | नमस्ते दोस्तों, आज हम बात करेंगे नई तकनीक के बारे में |

Each is 3.69s (3s real prompt + 0.69s AR). Still collapses after ~52 AR steps at this early epoch. Give them a listen — the key question is whether the voice quality and style transfers from the prompt to the new text.


it is not saying this text check once code once


The user is saying the generated audio is not speaking the new text. This makes sense — the model is likely just continuing the audio pattern from the prompt without actually conditioning on the new text.

Let me think about why:

1. The predictor is a transformer decoder with cross-attention to text. But it was only trained for 45 epochs — the text conditioning might not be strong enough yet.

2. More fundamentally, the voice prompting approach feeds real audio embeddings as input. The predictor sees real audio context and generates a continuation. But the text cross-attention might not be powerful enough to override the audio context — it just continues the audio pattern regardless of text.

3. This is actually expected behavior at this stage.

During training, the model only ever saw aligned audio and text from the same utterance. When I feed it a mismatched audio prompt with new text during inference, it's encountering something it's never trained on. The model learned to predict the next audio frame given previous audio plus matching text, so when those two inputs conflict, it defaults to following the audio pattern. This is a fundamental limitation — the model can't actually do voice cloning with arbitrary new text because it was never trained to handle that scenario.

For this to work, I'd need the text encoder to control what's being said while the audio prompt controls the voice identity, but the JEPA predictor isn't designed that way. It just predicts the next embedding, and with the text cross-attention probably too weak at epoch 45, the text conditioning might not even be functioning properly. I should test whether the model actually uses text at all by comparing predictions with and without it. Interestingly, the samples the user praised were using the same text as the original audio, so the inputs were actually aligned in those cases.


You're right — the model isn't actually speaking the new text. Let me verify:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, numpy as np
from model_v5 import LeWMTTSv5

device = torch.device('cuda:0')
ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()

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

# Test: does the predictor actually use text?
# Generate with two different texts from the same start_emb
with torch.no_grad():
    text_emb_1 = model.text_encoder(text_to_tokens('नमस्ते'))
    text_emb_2 = model.text_encoder(text_to_tokens('भारत एक महान देश है'))

    start = model.start_emb.to(device)

    # Generate 20 steps with text 1
    cache1 = model.init_ar_cache()
    emb1 = start
    outputs1 = []
    for i in range(20):
        emb1, cache1 = model.predict_next_cached(emb1, text_emb_1, i, cache1)
        outputs1.append(emb1)

    # Generate 20 steps with text 2
    cache2 = model.init_ar_cache()
    emb2 = start
    outputs2 = []
    for i in range(20):
        emb2, cache2 = model.predict_next_cached(emb2, text_emb_2, i, cache2)
        outputs2.append(emb2)

    out1 = torch.cat(outputs1, dim=1)
    out2 = torch.cat(outputs2, dim=1)

    diff = (out1 - out2).abs().mean().item()
    cos_sim = torch.nn.functional.cosine_similarity(out1.reshape(-1, 256), out2.reshape(-1, 256), dim=-1).mean().item()

    print(f'Different text comparison:')
    print(f'  L1 diff: {diff:.4f}')
    print(f'  Cosine sim: {cos_sim:.4f}')
    print(f'  If cosine ~1.0, text conditioning is NOT working')
" 2>&1

Check if text conditioning actually works


Different text comparison:
  L1 diff: 0.0135
  Cosine sim: 1.0000
  If cosine ~1.0, text conditioning is NOT working
Shell cwd was reset to /home/ubuntu/pipecat


**Cosine similarity = 1.0** — the text conditioning is completely ignored. The predictor generates identical output regardless of text. It's just doing audio continuation.

This is why the prompted samples sounded good — they were continuing the same audio, not speaking new text. The predictor learned to predict next audio frame purely from audio context, never learned to use text.

This is a fundamental issue with the JEPA approach for TTS:
- The prediction task (predict next frame from previous frames) can be solved **without text** — consecutive audio frames are highly correlated
- The cross-attention to text has no gradient pressure to be useful
- The model takes the shortcut: just copy/extrapolate from audio context

To make text actually control the output, we need:
1. **Explicit text-audio alignment** (duration predictor, like FastSpeech/VITS)
2. **Or make text essential** — e.g., mask out audio context so the model MUST use text
3. **Or discrete token prediction** where text tokens directly map to audio tokens

This is why production TTS systems don't use pure next-frame prediction with text as optional context. They force text alignment.

Want me to try approach 2 — **randomly masking audio context during training** so the predictor is forced to rely on text? Or should we pivot to a fundamentally different architecture?


are u sure we dng wrong in traing ... or do we dng in infernce


The user is asking a good question - is the problem in training or inference?

Let me think about this:

**During training:**
- Input: [start, z[0], z[1], ..., z[T-2]]
- Target: [z[0], z[1], ..., z[T-1]]
- Text is provided via cross-attention

The model is supposed to predict the next audio embedding given previous audio embeddings + text. But the problem is: the next audio embedding is ALWAYS highly correlated with the previous one (only 13ms apart at 75Hz). So the model can achieve low loss by just copying/extrapolating from the audio context, completely ignoring text.

This is fundamentally a training issue - the audio signal is so predictable on its own that the model never learns to rely on text. Even in systems like VALL-E that work well, they predict discrete tokens where the next token isn't trivially predictable from just the previous one. With continuous embeddings sampled at 75Hz, the prediction becomes almost deterministic without needing the text at all.

The architecture itself seems fine - text embeddings are available through cross-attention during both training and inference via KV-cache. But the gradient signal from the prediction loss isn't flowing strongly enough through cross-attention to text because self-attention to previous audio embeddings already provides sufficient information. I should check whether the text cross-attention weights are actually attending to specific tokens or just staying uniform, and verify that the training setup isn't accidentally preventing text from being used - like through an incorrect mask or a disconnected cross-attention layer. Continuing to examine the gradient flow, I'm wondering if the self-attention mechanism alone could minimize the loss without properly leveraging the cross-attention to the text embeddings, which might explain why the model isn't learning the text conditioning effectively.

This points to a fundamental design issue: the architecture supports text conditioning, but the training objective doesn't actually require it. The next-frame prediction task can be solved using just the audio self-attention, so the model has no incentive to use the text. I could address this by randomly masking portions of the audio during training to force the model to rely on text, or by adding an information bottleneck that makes the audio path insufficient on its own, or by introducing a text-audio alignment loss to encourage meaningful cross-attention weights.

The most practical approach is actually simpler—randomly drop the audio input entirely during some percentage of training steps, similar to classifier-free guidance in diffusion models. This forces the model to learn to generate from text alone, while still allowing both modalities at inference time. Let me implement this approach.


Good question. Let me check if it's a training issue or inference issue.


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import LeWMTTSv5

device = torch.device('cuda:0')
ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()

# Check: does text affect TF prediction at all?
import json
manifest = json.load(open('processed_data_codec/manifest.json'))
codec_emb = torch.load(manifest[0]['emb_path'], weights_only=True).unsqueeze(0).to(device)

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

with torch.no_grad():
    z = model.proj_in(codec_emb.transpose(1, 2))
    start = model.start_emb.expand(1, -1, -1)
    input_emb = torch.cat([start, z[:, :-1]], dim=1)

    # TF with correct text
    text_emb_real = model.text_encoder(text_to_tokens(manifest[0]['text']))
    pred_real = model.predictor(input_emb, text_emb_real)

    # TF with completely different text
    text_emb_fake = model.text_encoder(text_to_tokens('यह एक बिलकुल अलग वाक्य है'))
    pred_fake = model.predictor(input_emb, text_emb_fake)

    # TF with random noise as text
    text_emb_noise = torch.randn_like(text_emb_real)
    pred_noise = model.predictor(input_emb, text_emb_noise)

    diff_text = (pred_real - pred_fake).abs().mean().item()
    diff_noise = (pred_real - pred_noise).abs().mean().item()
    cos_text = torch.nn.functional.cosine_similarity(
        pred_real.reshape(-1, 256), pred_fake.reshape(-1, 256), dim=-1).mean().item()
    cos_noise = torch.nn.functional.cosine_similarity(
        pred_real.reshape(-1, 256), pred_noise.reshape(-1, 256), dim=-1).mean().item()

    print(f'TF with real vs different text:')
    print(f'  L1 diff: {diff_text:.6f}, cosine: {cos_text:.6f}')
    print(f'TF with real vs random noise text:')
    print(f'  L1 diff: {diff_noise:.6f}, cosine: {cos_noise:.6f}')
    print()
    print(f'Prediction norm: {pred_real.norm(dim=-1).mean():.2f}')
    print()
    if cos_text > 0.999:
        print('CONCLUSION: Text is COMPLETELY ignored during training.')
        print('The predictor learned to predict next frame from audio alone.')
        print('At 75Hz, consecutive frames are so similar that text is unnecessary.')
    elif cos_text > 0.99:
        print('CONCLUSION: Text has very weak influence.')
    else:
        print('CONCLUSION: Text has some influence.')
" 2>&1

Check if text conditioning works even in TF mode


TF with real vs different text:
  L1 diff: 0.010700, cosine: 0.999992
TF with real vs random noise text:
  L1 diff: 0.234930, cosine: 0.997577

Prediction norm: 50.15

CONCLUSION: Text is COMPLETELY ignored during training.
The predictor learned to predict next frame from audio alone.
At 75Hz, consecutive frames are so similar that text is unnecessary.
Shell cwd was reset to /home/ubuntu/pipecat


**It's a training problem, not inference.** Even in TF mode, swapping the text makes zero difference (cosine=0.9999). The predictor completely ignores text because at 75Hz, it can predict the next frame just from the previous audio frame — text is unnecessary.

The fix: **force the model to use text** by randomly dropping audio context during training. This is called **classifier-free guidance (CFG)**. With some probability, we replace audio inputs with zeros, so the model MUST use text to predict.


kill $(ps aux | grep train_v5_pred | grep python | grep -v grep | awk '{print $2}') 2>/dev/null; sleep 5; nvidia-smi | grep MiB

Kill current training


| N/A   39C    P0             54W /  400W |       1MiB /  81920MiB |      0%      Default |


   115→            T = codec_emb.shape[2]
   116→            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]
   117→
   118→            # Text encoding (trainable)
   119→            text_emb = model.text_encoder(text_tokens, text_mask)
   120→
   121→            # Project with FROZEN proj_in — targets are stable
   122→            with torch.no_grad():
   123→                z = model.proj_in(codec_seq)  # [B, T, d_model]
   124→                target_emb = z.clone()  # no EMA, direct target
   125→
   126→            # Prepend start, shift right
   127→            start = model.start_emb.expand(B, -1, -1)
   128→            input_emb = torch.cat([start, z[:, :-1]], dim=1)
   129→
   130→            # Input noise
   131→            noise_progress = min(1.0, global_step / (total_steps * 0.3))
   132→            noise_level = args.input_noise_max * noise_progress
   133→            if noise_level > 0:
   134→                noise = torch.zeros_like(input_emb)
   135→                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
   136→                                            device=device) * noise_level
   137→                input_emb = input_emb + noise
   138→
   139→            # Masks
   140→            if codec_mask is not None:
   141→                start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
   142→                pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
   143→                loss_mask = codec_mask
   144→            else:
   145→                pred_mask = None
   146→                loss_mask = None
   147→
   148→            # Predict
   149→            predicted = model.predictor(input_emb, text_emb, pred_mask, text_mask)
   150→
   151→            # Prediction loss (MSE vs frozen proj_in targets)
   152→            if loss_mask is not None:
   153→                valid = (~loss_mask).unsqueeze(-1)
   154→                pred_loss = F.mse_loss(
   155→                    predicted * valid, target_emb * valid, reduction="sum"
   156→                ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   157→            else:
   158→                pred_loss = F.mse_loss(predicted, target_emb)
   159→
   160→            optimizer.zero_grad()
   161→            pred_loss.backward()
   162→            grad_norm = torch.nn.utils.clip_grad_norm_(
   163→                [p for p in model.parameters() if p.requires_grad], args.grad_clip
   164→            )

            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)

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


Edit

            # 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

            # 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


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

            # 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

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

            # 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")
    args = parser.parse_args()
    train(args)

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


    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()

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


    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 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")
    args = parser.parse_args()
    train(args)

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


Now launch with audio dropping at 20% — the model MUST learn to use text for those samples:


Bash

cd /home/ubuntu/lewm-tts && rm -rf output_v5_pred && nohup python3 -u train_v5_pred.py \
    --pretrained output_v5/checkpoints/epoch_0120.pt \
    --data_dir processed_data_codec \
    --output_dir output_v5_pred \
    --epochs 200 \
    --batch_size 256 \
    --lr 1e-4 \
    --input_noise_max 0.05 \
    --audio_drop_prob 0.2 \
    --max_codec_frames 900 \
    --num_workers 4 \
    --print_every 20 \
    --save_every 10 \
    --reinit \
    > train_v5_pred.log 2>&1 &
echo "PID: $!"

Launch predictor training with 20% audio dropout (CFG)


PID: 3755095


sleep 45 && tail -10 /home/ubuntu/lewm-tts/train_v5_pred.log

Check CFG training started


Device: cuda
Loaded pretrained from output_v5/checkpoints/epoch_0120.pt (epoch 120)
Reinitializing predictor + text encoder from scratch
Trainable params: 9.81M (predictor + text encoder + start_emb)
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 26 batches/epoch
  [1/200] step 20 | pred=11.2790 | grad=4.21 | 29.0s
Epoch 1/200 | pred=11.3237 | 37.5s
  New best (pred=11.3237)


Training started with 20% audio dropout. Now the model can't just rely on audio context — 20% of the time audio input is zeroed out and it MUST use text to predict.

This will make the loss higher (harder task) but should result in actual text conditioning. ~37s/epoch, 200 epochs ≈ 2 hours. I'll check with samples once it converges enough.


status


Check training progress


  Saved: output_v5_pred/checkpoints/epoch_0030.pt
  New best (pred=0.6235)
  [31/200] step 800 | pred=0.6298 | grad=0.60 | 28.3s
Epoch 31/200 | pred=0.6314 | 36.8s
  [32/200] step 820 | pred=0.5977 | grad=0.43 | 19.8s


Epoch 32, pred_loss=0.62. Compare to the no-dropout run which plateaued at 0.44 — the loss is higher as expected (harder task with 20% audio dropout). Still dropping.

About 1h45m left. Want me to generate samples at a checkpoint to test if text conditioning is working now?


samples


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

ckpt = torch.load('output_v5_pred/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTSv5(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()
print(f'Loaded epoch {ckpt[\"epoch\"]+1}, pred_loss={ckpt[\"loss\"]:.4f}')

codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5_cfg', exist_ok=True)

def save_wav(w, path, sr=24000):
    mx = np.abs(w).max()
    if mx > 0: w = w / mx * 0.95
    torchaudio.save(path, torch.from_numpy(w).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(w)/sr:.2f}s)')

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

# First check: does text conditioning work now?
with torch.no_grad():
    text_emb_1 = model.text_encoder(text_to_tokens('नमस्ते'))
    text_emb_2 = model.text_encoder(text_to_tokens('भारत एक महान देश है'))

    start = model.start_emb.to(device)
    cache1 = model.init_ar_cache()
    cache2 = model.init_ar_cache()
    e1, e2 = start, start
    o1, o2 = [], []
    for i in range(20):
        e1, cache1 = model.predict_next_cached(e1, text_emb_1, i, cache1)
        e2, cache2 = model.predict_next_cached(e2, text_emb_2, i, cache2)
        o1.append(e1); o2.append(e2)

    out1 = torch.cat(o1, dim=1)
    out2 = torch.cat(o2, dim=1)
    cos = torch.nn.functional.cosine_similarity(out1.reshape(-1, 256), out2.reshape(-1, 256), dim=-1).mean().item()
    print(f'Text conditioning test: cosine={cos:.4f} (was 1.0000 before, lower=text matters)')

# Voice-prompted samples with new text
entry = manifest[0]
audio_path = entry['audio_path']
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)

with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    z_full = model.proj_in(codec_emb.transpose(1, 2))
    prompt_frames = int(3 * 75)
    prompt_z = z_full[:, :prompt_frames]

    # Save ground truth and prompt
    save_wav(wav_24k.squeeze().numpy(), 'samples_v5_cfg/gt.wav')

    new_texts = [
        ('original', entry['text']),
        ('new1', 'आज मौसम बहुत अच्छा है और हम सब खुश हैं'),
        ('new2', 'भारत एक महान देश है जहां अनेक भाषाएं बोली जाती हैं'),
    ]

    for name, text in new_texts:
        text_emb = model.text_encoder(text_to_tokens(text))

        all_embs = []
        cache = model.init_ar_cache()
        start = model.start_emb.to(device)
        out, cache = model.predict_next_cached(start, text_emb, 0, cache)
        all_embs.append(out)

        for i in range(prompt_frames - 1):
            out, cache = model.predict_next_cached(prompt_z[:, i:i+1], text_emb, i + 1, cache)
            all_embs.append(out)

        next_emb = prompt_z[:, -1:]
        ar_steps = 0
        for step in range(prompt_frames, prompt_frames + 375):
            out, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
            all_embs.append(out)
            next_emb = out
            ar_steps += 1
            if ar_steps > 50:
                norms = [e.norm().item() for e in all_embs[-20:]]
                if np.std(norms) < 0.01 * np.mean(norms): break

        real_part = prompt_z
        ar_part = torch.cat(all_embs[prompt_frames:], dim=1)
        combined = torch.cat([real_part, ar_part], dim=1)

        codec_out = model.proj_out(combined).transpose(1, 2)
        wav_out = codec.decoder(codec_out).squeeze().cpu().numpy()
        save_wav(wav_out, f'samples_v5_cfg/{name}.wav')
        print(f'  {name}: {text[:40]}... | AR={ar_steps} steps')

    # Pure synthesis (no prompt) with text only
    for name, text in [('synth_pure', 'नमस्ते दोस्तों आज हम बात करेंगे')]:
        text_emb = model.text_encoder(text_to_tokens(text))
        all_embs = [model.start_emb.to(device)]
        cache = model.init_ar_cache()
        next_emb, cache = model.predict_next_cached(all_embs[0], text_emb, 0, cache)
        all_embs.append(next_emb)
        for step in range(1, 750):
            next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
            all_embs.append(next_emb)
            if step > 50:
                norms = [e.norm().item() for e in all_embs[-20:]]
                if np.std(norms) < 0.01 * np.mean(norms): break
        pred_embs = torch.cat(all_embs, dim=1)
        codec_synth = model.proj_out(pred_embs).transpose(1, 2)
        save_wav(codec.decoder(codec_synth).squeeze().cpu().numpy(), f'samples_v5_cfg/{name}.wav')
        print(f'  {name}: {step+1} steps')

print('\\nAll in samples_v5_cfg/')
" 2>&1

Generate samples with CFG-trained model, test text conditioning


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loaded epoch 30, pred_loss=0.6235
Text conditioning test: cosine=0.9999 (was 1.0000 before, lower=text matters)
Saved: samples_v5_cfg/gt.wav (9.52s)
Saved: samples_v5_cfg/original.wav (4.03s)
  original: मेरे प्यारे देशवासियों, नमस्कार। मन की ब... | AR=77 steps
Saved: samples_v5_cfg/new1.wav (3.96s)
  new1: आज मौसम बहुत अच्छा है और हम सब खुश हैं... | AR=72 steps
Saved: samples_v5_cfg/new2.wav (3.96s)
  new2: भारत एक महान देश है जहां अनेक भाषाएं बोल... | AR=72 steps
Saved: samples_v5_cfg/synth_pure.wav (1.99s)
  synth_pure: 148 steps

All in samples_v5_cfg/
Shell cwd was reset to /home/ubuntu/pipecat


Cosine = 0.9999 — text conditioning still barely working at epoch 30. Too early. But some positives:

- AR continuation is longer now (72-77 steps vs 51 before)
- Pure synth produced 1.99s (148 steps) vs 1.5s before

Samples in `samples_v5_cfg/`:

| File | Description |
|------|-------------|
| `gt.wav` | Ground truth |
| `original.wav` | 3s prompt + AR, original text (4.03s) |
| `new1.wav` | 3s prompt + AR, new text (3.96s) |
| `new2.wav` | 3s prompt + AR, new text (3.96s) |
| `synth_pure.wav` | Pure text-only synthesis (1.99s) |

Listen to these, but text conditioning needs more training. The 20% audio dropout means only ~6 samples per batch are text-only — might need more epochs for the text pathway to strengthen. Training is at epoch 32/200, still has 170 epochs to go.


synth_pure.wav │ Pure text-only synthesis (1.99s - noise . and new1 and new2 also cmng up the same text as original