The user is asking two things:
1. Is the JEPA predictor actually predicting well? (pred_loss went from 0.10 to 0.073)
2. If prediction is good, we can swap out the mel decoder + Vocos for a neural codec like EnCodec/NanoCodec

This is a great insight! If the JEPA predictor is learning good representations and predictions, the bottleneck is just the mel→audio pipeline. Using a neural codec (which has its own high-quality decoder) could bypass the weak mel decoder entirely.

Let me check the prediction quality more carefully.


Great thinking! If the predictor is good, we just need a better audio backend. Let me check prediction quality.


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
import torch, numpy as np, json
from model import LeWMTTS

device = torch.device('cuda')
ckpt = torch.load('output_v4/checkpoints/epoch_0120.pt', map_location=device, weights_only=False)
model = LeWMTTS(ckpt['config']).to(device)
model.load_state_dict(ckpt['model'])
model.eval()

manifest = json.load(open('processed_data_100mel/manifest.json'))

# Test prediction quality on multiple samples
pred_losses = []
recon_losses = []
for i in range(20):
    mel = torch.load(manifest[i]['mel_path'], weights_only=True).unsqueeze(0).to(device)
    text = manifest[i]['text']
    text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).to(device)

    with torch.no_grad():
        text_emb = model.text_encoder(text_tokens)
        _, mu, _ = model.audio_encoder(mel)

        # EMA targets
        model.ema_audio_encoder.eval()
        _, mu_ema, _ = model.ema_audio_encoder(mel)
        target = mu_ema

        # Teacher-forced prediction: input=[start, mu[0],...,mu[T-2]], target=[ema[0],...,ema[T-1]]
        B = mu.shape[0]
        start = model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, mu[:, :-1]], dim=1)
        predicted = model.predictor(input_emb, text_emb)

        pred_mse = ((predicted - target)**2).mean().item()
        pred_cosine = torch.nn.functional.cosine_similarity(
            predicted.reshape(-1, predicted.shape[-1]),
            target.reshape(-1, target.shape[-1]),
            dim=-1
        ).mean().item()

        # Also check: decode predicted embeddings vs decode target embeddings
        mel_from_pred = model.mel_decoder(predicted)
        mel_from_target = model.mel_decoder(target)
        T = min(mel_from_pred.shape[2], mel_from_target.shape[2], mel.shape[2])

        pred_mel_l1 = (mel_from_pred[:,:,:T] - mel[:,:,:T]).abs().mean().item()
        target_mel_l1 = (mel_from_target[:,:,:T] - mel[:,:,:T]).abs().mean().item()

        pred_losses.append(pred_mse)

        if i < 5:
            print(f'Sample {i}: pred_MSE={pred_mse:.4f}, cosine_sim={pred_cosine:.4f}, '
                  f'mel_from_pred_L1={pred_mel_l1:.4f}, mel_from_target_L1={target_mel_l1:.4f}')

print(f'\\nAvg pred MSE over 20 samples: {np.mean(pred_losses):.4f}')
print(f'\\n--- Embedding space analysis ---')

# Check how similar consecutive embeddings are (predictor's job difficulty)
mel = torch.load(manifest[0]['mel_path'], weights_only=True).unsqueeze(0).to(device)
with torch.no_grad():
    _, mu, _ = model.audio_encoder(mel)
    model.ema_audio_encoder.eval()
    _, mu_ema, _ = model.ema_audio_encoder(mel)

    # Consecutive cosine similarity in EMA space
    cos_consec = torch.nn.functional.cosine_similarity(
        mu_ema[:, :-1].reshape(-1, mu_ema.shape[-1]),
        mu_ema[:, 1:].reshape(-1, mu_ema.shape[-1]),
        dim=-1
    )
    print(f'Consecutive EMA cosine sim: mean={cos_consec.mean():.4f}, std={cos_consec.std():.4f}')
    print(f'EMA embedding norm: mean={mu_ema.norm(dim=-1).mean():.2f}')
    print(f'Online embedding norm: mean={mu.norm(dim=-1).mean():.2f}')

    # Check prediction vs just copying previous embedding (baseline)
    start = model.start_emb.expand(1, -1, -1)
    input_emb = torch.cat([start, mu[:, :-1]], dim=1)
    predicted = model.predictor(input_emb, model.text_encoder(
        torch.tensor(list(manifest[0]['text'].encode('utf-8')), dtype=torch.long).unsqueeze(0).to(device)
    ))

    copy_baseline_mse = ((mu_ema[:, 1:] - mu_ema[:, :-1])**2).mean().item()
    pred_vs_target_mse = ((predicted - mu_ema)**2).mean().item()
    print(f'\\nCopy-previous baseline MSE: {copy_baseline_mse:.4f}')
    print(f'Predictor MSE: {pred_vs_target_mse:.4f}')
    print(f'Predictor is {copy_baseline_mse/pred_vs_target_mse:.1f}x better than copy baseline')
" 2>&1

Analyze JEPA prediction quality in detail


Sample 0: pred_MSE=0.0432, cosine_sim=0.9670, mel_from_pred_L1=0.6014, mel_from_target_L1=0.5700
Sample 1: pred_MSE=0.0407, cosine_sim=0.9687, mel_from_pred_L1=0.6015, mel_from_target_L1=0.5743
Sample 2: pred_MSE=0.0352, cosine_sim=0.9721, mel_from_pred_L1=0.5809, mel_from_target_L1=0.5468
Sample 3: pred_MSE=0.0484, cosine_sim=0.9626, mel_from_pred_L1=0.6460, mel_from_target_L1=0.6125
Sample 4: pred_MSE=0.0389, cosine_sim=0.9699, mel_from_pred_L1=0.6117, mel_from_target_L1=0.5845

Avg pred MSE over 20 samples: 0.0426

--- Embedding space analysis ---
Consecutive EMA cosine sim: mean=0.7887, std=0.2179
EMA embedding norm: mean=12.73
Online embedding norm: mean=13.60

Copy-previous baseline MSE: 0.2697
Predictor MSE: 0.0432
Predictor is 6.2x better than copy baseline
Shell cwd was reset to /home/ubuntu/pipecat


The JEPA predictor is working well:

- **Cosine similarity: 0.97** — predicted embeddings are very close to targets
- **6.2x better than copy baseline** — it's genuinely predicting, not just copying
- **Consecutive cosine sim: 0.79** — embeddings are distinct enough (not trivial)

But look at this: `mel_from_target_L1=0.57` — even decoding the **perfect EMA embeddings** gives L1=0.57. The mel decoder is the bottleneck, not the predictor.

You're exactly right — if we swap the mel decoder + Vocos for a proper neural codec, we keep the working JEPA predictor and get much better audio quality. The predictor predicts in embedding space, so we just need a codec whose embeddings we can predict.

Let me check what's available for a codec approach.


pip list 2>/dev/null | grep -iE "encodec|snac|dac|codec|nano"

Check installed audio codecs


Install EnCodec


audiocodecs                              0.0.1
dacvae                                   1.0.0
descript-audio-codec                     1.0.0
encodec                                  0.1.1
linacodec                                0.0.11
snac                                     1.2.1
SudachiDict-core                         20260116
SudachiPy                                0.6.10
torchcodec                               0.8.0
xcodec2                                  0.1.5


Requirement already satisfied: MarkupSafe>=2.0 in /home/ubuntu/.local/lib/python3.10/site-packages (from jinja2->torch->encodec) (3.0.3)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)


We have a lot of codecs available! Let me think about the best approach:

1. **EnCodec** - Facebook's codec, 24kHz, RVQ tokens (discrete). Well-known, good quality.
2. **SNAC** - Multi-scale codec, also discrete tokens.
3. **DAC (Descript Audio Codec)** - Good quality, discrete tokens.
4. **XCodec2** - Newer codec.

The key architectural question is: the current JEPA works in continuous embedding space (d_model=256 vectors). Neural codecs typically produce:
- **Continuous latents**: before VQ quantization - these are continuous embeddings we could predict

- **Discrete tokens**: after VQ - these require a different prediction approach (classification instead of MSE)

For JEPA, continuous embeddings are ideal. I could use a codec's encoder to get continuous embeddings, train the JEPA predictor on those, then decode through the codec's decoder—which is essentially what was originally intended with the mel decoder as a placeholder.

Looking at which codec works best, EnCodec at 24kHz produces embeddings at 75Hz with 320 samples per frame, and the continuous embeddings before quantization would be the right target for prediction.

Comparing the options: EnCodec gives 128-dim embeddings at 75Hz, DAC offers either 86Hz or 75Hz depending on sample rate with 1024-dim latents, and SNAC is more complex with multi-scale outputs. For JEPA I need good reconstruction quality, smooth embeddings that are predictable, and a reasonable frame rate. EnCodec's 75Hz is significantly higher than my current setup which gets around 21.5 embeddings per second, so that's a meaningful difference. The user also mentioned "nanocodec" which I should look into.

The core strategy would be straightforward: encode audio through a neural codec to get continuous latents, swap out the AudioEncoder and MelDecoder for the codec's encoder and decoder, and keep the JEPA predictor and TextEncoder but retrain them to work in the codec's embedding space. Since the codec's decoder is already trained on massive data and produces high-quality reconstructions, I just need the JEPA predictor to learn to forecast well in that latent space.

I'm leaning toward DAC since it's already available and performs well, though I should verify it handles 24kHz properly. The key advantage here is that I can freeze the codec entirely and only train the text encoder and predictor, which is much simpler than training custom audio encoders and decoders from scratch.

But there's a critical issue with vector quantization I need to address: the codec encoder produces continuous embeddings that get quantized before the decoder sees them, so if I predict continuous embeddings directly, the decoder might produce artifacts since it's never encountered un-quantized inputs during training. I could either quantize my predictions before decoding, fine-tune the decoder to handle continuous inputs at the risk of degrading quality, or find a codec that naturally works well with continuous embeddings.

Looking at EnCodec and DAC specifically, they sum embeddings from all RVQ codebook levels into a single continuous vector that the decoder processes, so if I can accurately predict that sum, the approach should work cleanly. I need to figure out which layer in these codecs to tap into for the prediction target, and I'm leaning toward EnCodec since it's the most established option among the available packages.


You have tons of codecs installed. Let me check EnCodec and DAC quality to pick the best one for JEPA.


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

# Test EnCodec
from encodec import EncodecModel
from encodec.utils import convert_audio

model = EncodecModel.encodec_model_24khz()
model.set_target_bandwidth(6.0)  # 6kbps
model.eval().cuda()

# Load a sample
import json
manifest = json.load(open('processed_data_100mel/manifest.json'))
wav_path = manifest[0]['audio_path']
wav, sr = torchaudio.load(wav_path)
print(f'Original: {wav.shape}, sr={sr}')

# Convert to 24kHz mono
wav = convert_audio(wav, sr, 24000, 1).cuda()
print(f'Converted: {wav.shape}')

with torch.no_grad():
    # Encode
    frames = model.encode(wav.unsqueeze(0))
    # frames is a list of (codes, scale) tuples
    codes = frames[0][0]  # [B, n_q, T]
    print(f'Codes shape: {codes.shape}')  # n_q = number of quantizers

    # Get continuous embeddings (before quantization)
    # EnCodec encoder output
    emb = model.encoder(wav.unsqueeze(0))
    print(f'Encoder continuous output: {emb.shape}')  # [B, channels, T_down]

    # Decode from codes
    decoded = model.decode(frames)
    print(f'Decoded: {decoded.shape}')

    # Also decode from continuous embeddings directly
    decoded_cont = model.decoder(emb)
    print(f'Decoded from continuous: {decoded_cont.shape}')

    # Quality comparison
    T = min(wav.shape[-1], decoded.shape[-1], decoded_cont.shape[-1])
    l1_quantized = (wav[:,:T] - decoded[0,:,:T]).abs().mean().item()
    l1_continuous = (wav[:,:T] - decoded_cont[0,:,:T]).abs().mean().item()
    print(f'\\nL1 from quantized codes: {l1_quantized:.6f}')
    print(f'L1 from continuous emb:  {l1_continuous:.6f}')

    print(f'\\nEmbedding rate: {emb.shape[2]} frames for {wav.shape[-1]/24000:.2f}s audio')
    print(f'Frame rate: {emb.shape[2] / (wav.shape[-1]/24000):.1f} Hz')
    print(f'Embedding dim: {emb.shape[1]}')
" 2>&1

Test EnCodec embedding structure and quality


/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)
Downloading: "https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th" to /home/ubuntu/.cache/torch/hub/checkpoints/encodec_24khz-d7cc33bc.th

  0%|          | 0.00/88.9M [00:00<?, ?B/s]
 41%|████      | 36.0M/88.9M [00:00<00:00, 377MB/s]
 96%|█████████▌| 85.5M/88.9M [00:00<00:00, 460MB/s]
100%|██████████| 88.9M/88.9M [00:00<00:00, 449MB/s]
Original: torch.Size([1, 228480]), sr=24000
Converted: torch.Size([1, 228480])
Codes shape: torch.Size([1, 8, 714])
Encoder continuous output: torch.Size([1, 128, 714])
Decoded: torch.Size([1, 1, 228480])
Decoded from continuous: torch.Size([1, 1, 228480])

L1 from quantized codes: 0.016688
L1 from continuous emb:  0.009797

Embedding rate: 714 frames for 9.52s audio
Frame rate: 75.0 Hz
Embedding dim: 128
Shell cwd was reset to /home/ubuntu/pipecat


Excellent results!

EnCodec 24kHz:
- **128-dim continuous embeddings** at 75 Hz
- **L1 from continuous: 0.0098** — near-transparent quality!
- **L1 from quantized: 0.017** — still very good
- 714 frames for 9.52s = 75 Hz frame rate

Compare to our current setup:
- 256-dim latents, 223 frames for 9.52s = ~23.4 Hz
- L1 from mel decoder: 0.53 on mel (then Vocos makes it worse)

The architecture is straightforward: freeze the EnCodec encoder and decoder, use them to get 128-dim embeddings at 75 Hz, then have the JEPA predictor work on those embeddings. The main tradeoff is that 75 Hz means 3.2x more frames to predict compared to our current setup, which could make autoregressive prediction slower, though I could add a downsampler to reduce the frame rate if needed. The decoder takes the predicted embeddings and reconstructs audio. I'm training the text encoder and JEPA predictor from scratch while keeping the EnCodec encoder and decoder frozen. For the frame rate, adding a stride-2 downsampler to reduce it to around 37.5 Hz seems like the right balance—it keeps temporal resolution reasonable while making the AR prediction more tractable.

Defining the downsampler module... Building the codec architecture...

The downsampler and upsampler handle the frequency conversion between the text encoder's 25Hz representation and the EnCodec's 75Hz codec space, so I can train the predictor to work with near-lossless audio embeddings instead of a weak decoder. For training, I'll precompute all the EnCodec embeddings upfront, then train the full pipeline end-to-end with MSE loss on the predicted embeddings. Now I'm weighing the loss functions needed to train the down/upsampler alongside JEPA's prediction objective — prediction loss on the latent space, reconstruction loss through the codec decoder, and an embedding loss to preserve information through the sampler layers. But this is getting complicated, so I'm reconsidering whether to skip the downsampler entirely and work directly at the codec's native 75Hz with 128 dimensions, accepting the slower autoregressive speed, or stick with downsampling to 25Hz and 256 dimensions to match the current setup but requiring more careful training of those conversion layers.

I'm going with the simpler approach — no downsampler, working at 75Hz with 128 dimensions projected to 256 for the JEPA predictor, then projecting back down before decoding. I'm setting up the pipeline to encode audio through the frozen codec, project the embeddings up for prediction, run them through JEPA, project back down, and decode through the frozen codec again. The losses are straightforward: prediction loss in the 256-dimensional space and a roundtrip loss to ensure the projections don't lose information. This is much cleaner, so I'm moving forward with implementation.

Actually, I'm reconsidering — I could work directly in 128 dimensions without any projection, but that might be too tight for the attention heads in the predictor. So I'm sticking with the 256-dimensional working space and adding projection layers to bridge between the codec's 128-dimensional output and the predictor's 256-dimensional input.

Setting up the projection layers...

Now I need to think through the full pipeline. The key bottleneck is preprocessing — I'll need to run EnCodec's encoder on all 7009 audio files upfront and cache the continuous embeddings to disk before training. Given the dataset size of about 17.88 hours, the disk footprint should be around 2.5 GB, which is totally workable. Let me start building out the preprocessing step.

Actually, I'm realizing I should keep this practical and modular. Rather than creating entirely new files, I can extend the existing codebase to support codec mode. But there's one more consideration: at 75 Hz, a 10-second clip would have 750 timesteps, which is significantly more than the current ~230 steps at 23 Hz. That could slow down inference, especially on CPU. I should add a stride-3 downsampler to bring it back down to around 25 Hz, keeping it closer to the current system's efficiency.

So my plan is to implement a codec-based approach with a downsampler that reduces the temporal resolution back to match what we're already using, then add an upsampler to reconstruct the full resolution. I'll create the preprocessing script to extract embeddings, build the model with the codec components and downsamplers, and keep the training straightforward with just prediction and reconstruction losses.


EnCodec is near-transparent:
- **Continuous embeddings**: 128-dim at 75Hz
- **L1 from continuous: 0.0098** vs our mel decoder's **0.53** — 54x better!

The JEPA predictor is solid (cosine_sim=0.97, 6.2x better than copy baseline). The weak mel decoder is the only problem. Swapping to EnCodec as the audio backend is the right move.

Here's the plan for v5:
- **EnCodec encoder/decoder**: frozen, handles audio ↔ embeddings
- **Stride-3 downsampler**: 75Hz → 25Hz (matches current AR speed)
- **Stride-3 upsampler**: 25Hz → 75Hz (back to codec rate)
- **JEPA Predictor + TextEncoder**: retrained in codec embedding space
- No mel decoder, no Vocos, no reconstruction loss headaches

Let me implement it.


ToolSearch

Question: select:Write,Edit,TaskCreate


Question: select:Write,Edit,TaskCreate


Let me create tasks and start implementing the v5 codec-based JEPA TTS.


TaskCreate

Preprocess audio to EnCodec embeddings

Create preprocess_codec.py to extract continuous EnCodec embeddings (128d, 75Hz) for all 7009 audio files and save to disk


Build v5 codec-based JEPA model

Create model_v5.py with: frozen EnCodec encoder/decoder, stride-3 down/upsampler (75Hz→25Hz→75Hz), TextEncoder, JEPA Predictor with EMA, learnable start_emb


Create v5 dataset and training loop

Create dataset_v5.py for codec embeddings and train_v5.py with simplified training (prediction loss + codec roundtrip loss)


Create v5 inference script

Create inference_v5.py: text → predictor → upsample → EnCodec decoder → waveform


Preprocess data and start v5 training

Run preprocessing, then launch 200-epoch training on A100 with optimal batch size


Task #6 created successfully: Preprocess audio to EnCodec embeddings

Preprocess audio to EnCodec embeddings


Task #7 created successfully: Preprocess data and start v5 training

Preprocess data and start v5 training


Task #8 created successfully: Create v5 dataset and training loop

Create v5 dataset and training loop


Task #9 created successfully: Create v5 inference script

Create v5 inference script


Task #10 created successfully: Build v5 codec-based JEPA model

Build v5 codec-based JEPA model


ToolSearch

Question: select:TaskUpdate


Question: select:TaskUpdate


Updated task #6 status


Write

"""
Preprocess audio files → EnCodec continuous embeddings.
Saves 128-dim embeddings at 75Hz for each audio file.
"""

import os
import json
import torch
import torchaudio
import argparse
from pathlib import Path
from encodec import EncodecModel
from encodec.utils import convert_audio

def preprocess(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

    # Load original manifest
    with open(args.manifest, "r") as f:
        manifest = json.load(f)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    emb_dir = out_dir / "embeddings"
    emb_dir.mkdir(exist_ok=True)

    new_manifest = []
    skipped = 0

    for i, entry in enumerate(manifest):
        audio_path = entry["audio_path"]
        if not os.path.exists(audio_path):
            skipped += 1
            continue

        try:
            wav, sr = torchaudio.load(audio_path)
            wav = convert_audio(wav, sr, 24000, 1).to(device)

            with torch.no_grad():
                # Get continuous embeddings (before VQ)
                emb = codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
                emb = emb.squeeze(0).cpu()  # [128, T]

            # Save embedding
            emb_path = str(emb_dir / f"{i:06d}.pt")
            torch.save(emb, emb_path)

            new_manifest.append({
                "emb_path": emb_path,
                "text": entry["text"],
                "audio_path": audio_path,
                "emb_frames": emb.shape[1],
                "duration": wav.shape[-1] / 24000,
            })

            if (i + 1) % 500 == 0:
                print(f"  Processed {i+1}/{len(manifest)} files")

        except Exception as e:
            print(f"  Error on {audio_path}: {e}")
            skipped += 1

    # Save new manifest
    manifest_path = str(out_dir / "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(new_manifest, f, indent=2, ensure_ascii=False)

    print(f"\nDone: {len(new_manifest)} files processed, {skipped} skipped")
    print(f"Manifest: {manifest_path}")
    print(f"Embeddings: {emb_dir}")

    # Stats
    total_frames = sum(e["emb_frames"] for e in new_manifest)
    total_dur = sum(e["duration"] for e in new_manifest)
    print(f"Total: {total_dur/3600:.2f}h, {total_frames} frames, avg {total_frames/len(new_manifest):.0f} frames/sample")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", default="processed_data_100mel/manifest.json")
    parser.add_argument("--output_dir", default="processed_data_codec")
    args = parser.parse_args()
    preprocess(args)


File created successfully at: /home/ubuntu/lewm-tts/preprocess_codec.py

"""
Preprocess audio files → EnCodec continuous embeddings.
Saves 128-dim embeddings at 75Hz for each audio file.
"""

import os
import json
import torch
import torchaudio
import argparse
from pathlib import Path
from encodec import EncodecModel
from encodec.utils import convert_audio

def preprocess(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

    # Load original manifest
    with open(args.manifest, "r") as f:
        manifest = json.load(f)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    emb_dir = out_dir / "embeddings"
    emb_dir.mkdir(exist_ok=True)

    new_manifest = []
    skipped = 0

    for i, entry in enumerate(manifest):
        audio_path = entry["audio_path"]
        if not os.path.exists(audio_path):
            skipped += 1
            continue

        try:
            wav, sr = torchaudio.load(audio_path)
            wav = convert_audio(wav, sr, 24000, 1).to(device)

            with torch.no_grad():
                # Get continuous embeddings (before VQ)
                emb = codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
                emb = emb.squeeze(0).cpu()  # [128, T]

            # Save embedding
            emb_path = str(emb_dir / f"{i:06d}.pt")
            torch.save(emb, emb_path)

            new_manifest.append({
                "emb_path": emb_path,
                "text": entry["text"],
                "audio_path": audio_path,
                "emb_frames": emb.shape[1],
                "duration": wav.shape[-1] / 24000,
            })

            if (i + 1) % 500 == 0:
                print(f"  Processed {i+1}/{len(manifest)} files")

        except Exception as e:
            print(f"  Error on {audio_path}: {e}")
            skipped += 1

    # Save new manifest
    manifest_path = str(out_dir / "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(new_manifest, f, indent=2, ensure_ascii=False)

    print(f"\nDone: {len(new_manifest)} files processed, {skipped} skipped")
    print(f"Manifest: {manifest_path}")
    print(f"Embeddings: {emb_dir}")

    # Stats
    total_frames = sum(e["emb_frames"] for e in new_manifest)
    total_dur = sum(e["duration"] for e in new_manifest)
    print(f"Total: {total_dur/3600:.2f}h, {total_frames} frames, avg {total_frames/len(new_manifest):.0f} frames/sample")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", default="processed_data_100mel/manifest.json")
    parser.add_argument("--output_dir", default="processed_data_codec")
    args = parser.parse_args()
    preprocess(args)


Now the v5 model with codec backend:


Updated task #10 status


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

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - CodecDownsampler: 75Hz → 25Hz, 128d → 256d (trainable)
  - CodecUpsampler: 25Hz → 75Hz, 256d → 128d (trainable)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on downsampled embeddings
"""

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

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

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

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

# ─── Codec Down/Upsampler ──────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class CodecDownsampler(nn.Module):
    """Downsample codec embeddings: 75Hz×128d → 25Hz×d_model."""
    def __init__(self, codec_dim=128, d_model=256, stride=3):
        super().__init__()
        self.stride = stride
        self.net = nn.Sequential(
            nn.Conv1d(codec_dim, d_model, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            nn.Conv1d(d_model, d_model, kernel_size=2 * stride + 1, stride=stride, padding=stride),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=3),
        )
        # Transformer layers for contextual understanding
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=4, dim_feedforward=d_model * 4,
            dropout=0.1, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

    def forward(self, codec_emb, mask=None):
        """codec_emb: [B, 128, T_75Hz] → [B, T_25Hz, d_model]"""
        x = self.net(codec_emb)  # [B, d_model, T_25Hz]
        x = x.transpose(1, 2)   # [B, T_25Hz, d_model]
        if mask is not None:
            mask = mask[:, ::self.stride][:, :x.shape[1]]
        x = self.pos_embed(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return x, mask

class CodecUpsampler(nn.Module):
    """Upsample predicted embeddings: 25Hz×d_model → 75Hz×128d."""
    def __init__(self, d_model=256, codec_dim=128, stride=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=2 * stride, stride=stride, padding=stride // 2),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            ResConvBlock(d_model, kernel_size=3),
            nn.Conv1d(d_model, codec_dim, kernel_size=7, padding=3),
        )

    def forward(self, z, target_len=None):
        """z: [B, T_25Hz, d_model] → [B, 128, T_75Hz]"""
        x = z.transpose(1, 2)  # [B, d_model, T_25Hz]
        x = self.net(x)        # [B, 128, T_75Hz_approx]
        if target_len is not None:
            if x.shape[2] > target_len:
                x = x[:, :, :target_len]
            elif x.shape[2] < target_len:
                x = F.pad(x, (0, target_len - x.shape[2]))
        return x

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

class JEPAPredictor(nn.Module):
    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(
            x, text_emb, tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,
        )
        return self.output_proj(predicted)

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

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

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

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        x = self.audio_input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))
        return self.output_proj(x), cache

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

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

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.

    Training flow:
        codec_emb → downsample → JEPA predict → loss vs EMA targets
        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

    Inference flow:
        text → predictor (AR) → upsample → EnCodec decoder → waveform
    """

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

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.downsampler = CodecDownsampler(
            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,
        )
        self.upsampler = CodecUpsampler(
            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

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

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

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

        self.config = config

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

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

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

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

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

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

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

        # ─── Loss 2: Roundtrip (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

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

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

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

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


File created successfully at: /home/ubuntu/lewm-tts/model_v5.py

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

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - CodecDownsampler: 75Hz → 25Hz, 128d → 256d (trainable)
  - CodecUpsampler: 25Hz → 75Hz, 256d → 128d (trainable)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on downsampled embeddings
"""

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

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

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

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

# ─── Codec Down/Upsampler ──────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class CodecDownsampler(nn.Module):
    """Downsample codec embeddings: 75Hz×128d → 25Hz×d_model."""
    def __init__(self, codec_dim=128, d_model=256, stride=3):
        super().__init__()
        self.stride = stride
        self.net = nn.Sequential(
            nn.Conv1d(codec_dim, d_model, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            nn.Conv1d(d_model, d_model, kernel_size=2 * stride + 1, stride=stride, padding=stride),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=3),
        )
        # Transformer layers for contextual understanding
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=4, dim_feedforward=d_model * 4,
            dropout=0.1, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

    def forward(self, codec_emb, mask=None):
        """codec_emb: [B, 128, T_75Hz] → [B, T_25Hz, d_model]"""
        x = self.net(codec_emb)  # [B, d_model, T_25Hz]
        x = x.transpose(1, 2)   # [B, T_25Hz, d_model]
        if mask is not None:
            mask = mask[:, ::self.stride][:, :x.shape[1]]
        x = self.pos_embed(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return x, mask

class CodecUpsampler(nn.Module):
    """Upsample predicted embeddings: 25Hz×d_model → 75Hz×128d."""
    def __init__(self, d_model=256, codec_dim=128, stride=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=2 * stride, stride=stride, padding=stride // 2),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            ResConvBlock(d_model, kernel_size=3),
            nn.Conv1d(d_model, codec_dim, kernel_size=7, padding=3),
        )

    def forward(self, z, target_len=None):
        """z: [B, T_25Hz, d_model] → [B, 128, T_75Hz]"""
        x = z.transpose(1, 2)  # [B, d_model, T_25Hz]
        x = self.net(x)        # [B, 128, T_75Hz_approx]
        if target_len is not None:
            if x.shape[2] > target_len:
                x = x[:, :, :target_len]
            elif x.shape[2] < target_len:
                x = F.pad(x, (0, target_len - x.shape[2]))
        return x

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

class JEPAPredictor(nn.Module):
    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(
            x, text_emb, tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,
        )
        return self.output_proj(predicted)

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

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

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

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        x = self.audio_input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))
        return self.output_proj(x), cache

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

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

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.

    Training flow:
        codec_emb → downsample → JEPA predict → loss vs EMA targets
        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

    Inference flow:
        text → predictor (AR) → upsample → EnCodec decoder → waveform
    """

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

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.downsampler = CodecDownsampler(
            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,
        )
        self.upsampler = CodecUpsampler(
            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

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

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

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

        self.config = config

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

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

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

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

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

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

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

        # ─── Loss 2: Roundtrip (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

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

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

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

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


Updated task #8 status


"""
Dataset for codec-based JEPA TTS v5.
Loads precomputed EnCodec embeddings + text.
"""

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

class CodecDataset(Dataset):
    def __init__(self, manifest_path, max_codec_frames=900):
        with open(manifest_path, "r") as f:
            self.manifest = json.load(f)
        self.max_codec_frames = max_codec_frames
        # Filter out too-long samples
        self.manifest = [e for e in self.manifest if e["emb_frames"] <= max_codec_frames]
        print(f"CodecDataset: {len(self.manifest)} samples (max {max_codec_frames} frames)")

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

    def __getitem__(self, idx):
        entry = self.manifest[idx]
        emb = torch.load(entry["emb_path"], weights_only=True)  # [128, T]
        text_tokens = torch.tensor(list(entry["text"].encode("utf-8")), dtype=torch.long)
        return {"codec_emb": emb, "text_tokens": text_tokens}

def collate_fn(batch):
    codec_embs = [b["codec_emb"] for b in batch]
    text_tokens = [b["text_tokens"] for b in batch]

    # Pad codec embeddings
    codec_dim = codec_embs[0].shape[0]
    max_codec_len = max(e.shape[1] for e in codec_embs)
    padded_codec = torch.zeros(len(batch), codec_dim, max_codec_len)
    codec_mask = torch.ones(len(batch), max_codec_len, dtype=torch.bool)
    for i, e in enumerate(codec_embs):
        T = e.shape[1]
        padded_codec[i, :, :T] = e
        codec_mask[i, :T] = False

    # Pad text
    max_text_len = max(t.shape[0] for t in text_tokens)
    padded_text = torch.zeros(len(batch), max_text_len, dtype=torch.long)
    text_mask = torch.ones(len(batch), max_text_len, dtype=torch.bool)
    for i, t in enumerate(text_tokens):
        padded_text[i, :t.shape[0]] = t
        text_mask[i, :t.shape[0]] = False

    return {
        "codec_emb": padded_codec,
        "text_tokens": padded_text,
        "codec_mask": codec_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=64, num_workers=4, max_codec_frames=900):
    dataset = CodecDataset(manifest_path, max_codec_frames=max_codec_frames)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=True,
        num_workers=num_workers, collate_fn=collate_fn,
        pin_memory=True, drop_last=True,
    )
    return loader, dataset


File created successfully at: /home/ubuntu/lewm-tts/dataset_v5.py

"""
Dataset for codec-based JEPA TTS v5.
Loads precomputed EnCodec embeddings + text.
"""

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

class CodecDataset(Dataset):
    def __init__(self, manifest_path, max_codec_frames=900):
        with open(manifest_path, "r") as f:
            self.manifest = json.load(f)
        self.max_codec_frames = max_codec_frames
        # Filter out too-long samples
        self.manifest = [e for e in self.manifest if e["emb_frames"] <= max_codec_frames]
        print(f"CodecDataset: {len(self.manifest)} samples (max {max_codec_frames} frames)")

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

    def __getitem__(self, idx):
        entry = self.manifest[idx]
        emb = torch.load(entry["emb_path"], weights_only=True)  # [128, T]
        text_tokens = torch.tensor(list(entry["text"].encode("utf-8")), dtype=torch.long)
        return {"codec_emb": emb, "text_tokens": text_tokens}

def collate_fn(batch):
    codec_embs = [b["codec_emb"] for b in batch]
    text_tokens = [b["text_tokens"] for b in batch]

    # Pad codec embeddings
    codec_dim = codec_embs[0].shape[0]
    max_codec_len = max(e.shape[1] for e in codec_embs)
    padded_codec = torch.zeros(len(batch), codec_dim, max_codec_len)
    codec_mask = torch.ones(len(batch), max_codec_len, dtype=torch.bool)
    for i, e in enumerate(codec_embs):
        T = e.shape[1]
        padded_codec[i, :, :T] = e
        codec_mask[i, :T] = False

    # Pad text
    max_text_len = max(t.shape[0] for t in text_tokens)
    padded_text = torch.zeros(len(batch), max_text_len, dtype=torch.long)
    text_mask = torch.ones(len(batch), max_text_len, dtype=torch.bool)
    for i, t in enumerate(text_tokens):
        padded_text[i, :t.shape[0]] = t
        text_mask[i, :t.shape[0]] = False

    return {
        "codec_emb": padded_codec,
        "text_tokens": padded_text,
        "codec_mask": codec_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=64, num_workers=4, max_codec_frames=900):
    dataset = CodecDataset(manifest_path, max_codec_frames=max_codec_frames)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=True,
        num_workers=num_workers, collate_fn=collate_fn,
        pin_memory=True, drop_last=True,
    )
    return loader, dataset


"""
Training loop for LeWM TTS v5 (codec-based JEPA).
"""

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

from model_v5 import 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}")

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "codec_dim": 128,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "ds_stride": args.ds_stride,
        "pred_weight": args.pred_weight,
        "roundtrip_weight": args.roundtrip_weight,
        "ema_decay": args.ema_decay,
        "input_noise": 0.0,
    }
    model, config = build_model(config)
    model = model.to(device)

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

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

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

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

    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)

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

    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_rt = 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)

            # Input noise annealing
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            losses = model(codec_emb, text_tokens, codec_mask, text_mask)

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            rt_loss = losses["roundtrip_loss"]

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()
            model.update_ema()

            epoch_loss += total_loss.item()
            epoch_pred += pred_loss.item()
            epoch_rt += rt_loss.item()
            global_step += 1

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

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

        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred / n_batches
        avg_rt = epoch_rt / n_batches
        epoch_time = time.time() - t0

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

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
            torch.save({
                "epoch": epoch,
                "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config,
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved checkpoint: {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 (loss={avg_loss:.4f})")

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5")
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--ds_stride", type=int, default=3)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--roundtrip_weight", type=float, default=1.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--input_noise_max", type=float, default=0.1)
    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=64)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=50)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args()
    train(args)


File created successfully at: /home/ubuntu/lewm-tts/train_v5.py

"""
Training loop for LeWM TTS v5 (codec-based JEPA).
"""

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

from model_v5 import 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}")

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "codec_dim": 128,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "ds_stride": args.ds_stride,
        "pred_weight": args.pred_weight,
        "roundtrip_weight": args.roundtrip_weight,
        "ema_decay": args.ema_decay,
        "input_noise": 0.0,
    }
    model, config = build_model(config)
    model = model.to(device)

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

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

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

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

    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)

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

    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_rt = 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)

            # Input noise annealing
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            losses = model(codec_emb, text_tokens, codec_mask, text_mask)

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            rt_loss = losses["roundtrip_loss"]

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()
            model.update_ema()

            epoch_loss += total_loss.item()
            epoch_pred += pred_loss.item()
            epoch_rt += rt_loss.item()
            global_step += 1

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

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

        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred / n_batches
        avg_rt = epoch_rt / n_batches
        epoch_time = time.time() - t0

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

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
            torch.save({
                "epoch": epoch,
                "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config,
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved checkpoint: {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 (loss={avg_loss:.4f})")

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5")
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--ds_stride", type=int, default=3)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--roundtrip_weight", type=float, default=1.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--input_noise_max", type=float, default=0.1)
    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=64)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=50)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args()
    train(args)


Updated task #9 status


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

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

from model_v5 import LeWMTTSv5

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

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

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

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

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

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

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

        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

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

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

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

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

        # Concat predicted embeddings and upsample to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

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

        return waveform, 24000

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

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

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)
        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

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

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

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

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

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

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

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

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

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

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

    tts = LeWMTTSv5Inference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/lewm-tts/inference_v5.py

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

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

from model_v5 import LeWMTTSv5

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

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

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

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

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

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

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

        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

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

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

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

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

        # Concat predicted embeddings and upsample to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

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

        return waveform, 24000

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

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

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)
        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

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

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

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

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

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

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

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

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

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

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

    tts = LeWMTTSv5Inference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


Updated task #7 status


Now let me preprocess the data and start training. First, preprocessing:


Bash

cd /home/ubuntu/lewm-tts && python3 -u preprocess_codec.py 2>&1

Extract EnCodec embeddings for all audio files


/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)
  Processed 500/7009 files
  Processed 1000/7009 files
  Processed 1500/7009 files
  Processed 2000/7009 files
  Processed 2500/7009 files
  Processed 3000/7009 files
  Processed 3500/7009 files
  Processed 4000/7009 files
  Processed 4500/7009 files
  Processed 5000/7009 files
  Processed 5500/7009 files
  Processed 6000/7009 files
  Processed 6500/7009 files
  Processed 7000/7009 files

Done: 7009 files processed, 0 skipped
Manifest: processed_data_codec/manifest.json
Embeddings: processed_data_codec/embeddings
Total: 17.88h, 4829533 frames, avg 689 frames/sample
Shell cwd was reset to /home/ubuntu/pipecat


All 7009 files preprocessed. Average 689 frames/sample at 75Hz → after stride-3 downsampling = ~230 frames/sample (same as v4). Now let's find the right batch size and start training:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader

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

loader, dataset = build_dataloader('processed_data_codec/manifest.json', batch_size=64)
batch = next(iter(loader))

# Test forward pass
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)

model.train()
losses = model(codec_emb, text_tokens, codec_mask, text_mask)
print(f'batch_size=64: pred={losses[\"prediction_loss\"].item():.4f}, rt={losses[\"roundtrip_loss\"].item():.4f}')
print(f'GPU mem: {torch.cuda.max_memory_allocated()/1e9:.1f}GB')

torch.cuda.reset_peak_memory_stats()
" 2>&1

Test v5 model forward pass and memory


Error: Exit code 1
LeWM TTS v5: 14.80M trainable parameters
CodecDataset: 6879 samples (max 900 frames)
/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py:515: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)
  output = torch._nested_tensor_from_mask(
Traceback (most recent call last):
  File "<string>", line 20, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/lewm-tts/model_v5.py", line 348, in forward
    predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/lewm-tts/model_v5.py", line 142, in forward
    predicted = self.transformer(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 628, in forward
    output = mod(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 1131, in forward
    x = self.norm3(x + self._ff_block(x))
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 1176, in _ff_block
    x = self.linear2(self.dropout(self.activation(self.linear1(x))))
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/dropout.py", line 73, in forward
    return F.dropout(input, self.p, self.training, self.inplace)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/functional.py", line 1418, in dropout
    _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 76.00 MiB. GPU 0 has a total capacity of 79.25 GiB of which 21.19 MiB is free. Process 3652358 has 74.20 GiB memory in use. Including non-PyTorch memory, this process has 5.01 GiB memory in use. Of the allocated memory 4.42 GiB is allocated by PyTorch, and 105.97 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


The v4 training is using 74GB. Let me kill it first since we're moving to v5.


Check v4 training progress before killing


  [160/200] step 3360 | loss=1.9416 pred=0.1026 kl=0.4085 recon=0.5661 spec=0.6578 | grad=2.80 | 41.7s
Epoch 160/200 | avg_loss=1.9206 pred=0.1010 kl=0.4069 recon=0.5620 spec=0.6573 | 41.8s
  Saved checkpoint: output_v4/checkpoints/epoch_0160.pt
  [161/200] step 3380 | loss=1.9338 pred=0.1033 kl=0.4146 recon=0.5611 spec=0.6390 | grad=1.66 | 37.8s
Epoch 161/200 | avg_loss=1.9382 pred=0.1027 kl=0.4120 recon=0.5623 spec=0.6560 | 39.2s


V4 is at epoch 161/200. Recon loss at 0.56 — barely improving. Let me kill it and start v5.


kill 3652358 && sleep 2 && nvidia-smi | grep -E "MiB|%"

Kill v4 training and check GPU is free


| N/A   40C    P0             65W /  400W |   75997MiB /  81920MiB |      0%      Default |


Wait for GPU memory to free


| N/A   36C    P0             53W /  400W |       1MiB /  81920MiB |      0%      Default |


GPU is free. Now let's find the right batch size for v5:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader

device = torch.device('cuda')

for bs in [64, 128, 256, 384]:
    torch.cuda.reset_peak_memory_stats()
    try:
        model, config = build_model()
        model = model.to(device)
        model.train()

        loader, dataset = build_dataloader('processed_data_codec/manifest.json', batch_size=bs)
        batch = next(iter(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)

        losses = model(codec_emb, text_tokens, codec_mask, text_mask)
        losses['total_loss'].backward()

        mem = torch.cuda.max_memory_allocated()/1e9
        print(f'batch_size={bs}: OK, GPU={mem:.1f}GB, pred={losses[\"prediction_loss\"].item():.4f}, rt={losses[\"roundtrip_loss\"].item():.4f}')

        del model, loader, dataset, batch, losses
        torch.cuda.empty_cache()
    except torch.cuda.OutOfMemoryError:
        print(f'batch_size={bs}: OOM')
        torch.cuda.empty_cache()
        break
" 2>&1 | grep -E "batch_size|OOM"

Find optimal batch size for v5 on A100