hi


Hi! How can I help you today?


'/home/ubuntu/lewm-tts' Yeah, I want you to recall your memory, like we tried with joint and mudding layers, something we did with jpeg chipa, and we couldn't get those good results after running 200 epochs. You can see we have the v3c version. I don't know why we didn't get the good results. I want you to go and thoroughly check what we did wrong and what will be done wrong. I am very confident that you can get good results with the jepa models.


Let me check my memory files first, then explore the project directory to understand what's been done.


ToolSearch

Question: select:Read,Glob,Grep,Agent,Bash


Question: select:Read,Glob,Grep,Agent,Bash


Let me check my memory files first, then explore the project directory.


Let me check my memory and then thoroughly explore the project.


Bash

List memory files


List project directory


Error: File does not exist. Note: your current working directory is /home/ubuntu/pipecat.


dataset.py
inference_engine.py
inference.py
model.py
output
output_v3
output_v3b
output_v3c
preprocess.py
processed_data
processed_data_100mel
__pycache__
README.md
run_v1.sh
sample_ar_prompted_v3b.wav
sample_ar_prompted.wav
sample_ar_v3c.wav
sample_gt_0.wav
sample_gt_100.wav
sample_gt_500.wav
sample_gt_v3b.wav
sample_recon_ep43_0.wav
sample_recon_ep43_100.wav
sample_recon_ep43_500.wav
sample_reconstruct_v3.wav
sample_recon_v3b_0.wav
sample_recon_v3b_100.wav
sample_recon_v3b_500.wav
sample_recon_v3b.wav
sample_recon_v3c_0.wav
sample_recon_v3c_100.wav
sample_recon_v3c_500.wav
samples_v3c
sample_synth_ep43.wav
sample_synth_v3b.wav
sample_synth_v3.wav
sample_tf_0.wav
sample_tf_100.wav
sample_tf_500.wav
sample_tf_v3b_0.wav
sample_tf_v3b_100.wav
sample_tf_v3c_0.wav
test_gt.wav
test_noisy_same_l1.wav
test_our_mel_vocos.wav
test_recon_v3c.wav
test_vocos_roundtrip.wav
train.py
train_v3b.log
train_v3c.log
train_v3.log
v2


(Bash completed with no output)


     1→# LeWM TTS
     2→
     3→JEPA-based Text-to-Speech system adapted from [LeWorldModel](https://le-wm.github.io/).
     4→
     5→## Architecture
     6→
     7→A 24M parameter model with 4 core components:
     8→
     9→```
    10→Text (Hindi/UTF-8 bytes) --> TextEncoder --> text embeddings ---+
    11→                                                                |--> JEPAPredictor --> predicted embeddings
    12→Audio (mel spectrogram)  --> AudioEncoder --> audio embeddings --+           |
    13→                                                                            v
    14→                                                                       MelDecoder --> mel spectrogram --> Vocos --> waveform
    15→```
    16→
    17→| Component | Description |
    18→|-----------|-------------|
    19→| **TextEncoder** | Byte-level character transformer (4 layers, 256d, 4 heads). Encodes Hindi text via UTF-8 bytes — no tokenizer needed. |
    20→| **AudioEncoder** | 1D CNN (stride-4 downsample) + Transformer (4 layers). Outputs mu/logvar for Gaussian latent space with reparameterization. |
    21→| **JEPAPredictor** | Transformer decoder (6 layers) with causal self-attention + cross-attention to text. Predicts next audio embedding autoregressively. Has KV-cache for fast inference. |
    22→| **MelDecoder** | Linear projection to 512d, 2x transposed convolution (4x upsample), residual conv blocks (ResConvBlock), and refinement layers. Outputs 100-mel spectrogram. |
    23→| **Vocoder** | [Vocos](https://github.com/gemelo-ai/vocos) `charactr/vocos-mel-24khz` — directly decodes 100-mel to 24kHz waveform. No Griffin-Lim. |
    24→
    25→### Speaker Embedding (planned)
    26→Built-in support for multi-speaker via learned embedding table added to text embeddings. Not yet trained — target is 10-20 speaker multi-speaker TTS.
    27→
    28→## Training Losses
    29→
    30→4 losses, all active during training:
    31→
    32→1. **Prediction Loss (MSE)** — next-embedding prediction: `MSE(predicted[t], target[t+1])`
    33→2. **KL Loss** — Gaussian regularizer: `KL(q(z|x) || N(0, I))` with annealing over first 20% of training
    34→3. **Reconstruction Loss (L1)** — mel decoder output vs ground truth mel spectrogram
    35→4. **Multi-Resolution Spectral Loss** — STFT at 3 resolutions (64/128/256 FFT sizes), combines spectral convergence + log magnitude L1. Forces sharp spectral detail instead of blurry averages.
    36→
    37→## Dataset
    38→
    39→- **Source**: 20 hours of studio-quality single-speaker Hindi audio
    40→- **Preprocessing**: silence-based segmentation into 3-12s chunks, mel extraction
    41→- **Processed**: 7,009 segments, 17.88 hours
    42→- **Mel config**: 100 mels, n_fft=1024, hop=256, sr=24000, f_max=None (full bandwidth), log-clamped at 1e-7 — matches Vocos `mel-24khz` exactly
    43→
    44→## Files
    45→
    46→| File | Purpose |
    47→|------|---------|
    48→| `model.py` | All model components: TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, MultiResolutionSpectralLoss, LeWMTTS |
    49→| `train.py` | Training loop with KL annealing, gradient clipping, TensorBoard logging, checkpointing |
    50→| `inference.py` | Single-request inference — AR synthesis, teacher-forced reconstruction |
    51→| `inference_engine.py` | Production batched inference engine — KV-cache, FP16, voice prompts, benchmarking |
    52→| `preprocess.py` | Data preprocessing — audio segmentation, mel extraction, manifest creation |
    53→| `dataset.py` | PyTorch dataset + dataloader with dynamic batching and padding |
    54→
    55→## Inference
    56→
    57→### Modes
    58→- **Synthesize**: Pure autoregressive generation from text. Requires audio prompt (voice cloning) for best results.
    59→- **Reconstruct**: Encode real mel -> decode back (tests encoder/decoder quality)
    60→- **Predict (teacher-forced)**: Encode mel -> predictor -> decode (tests full pipeline without AR drift)
    61→
    62→### Performance (A100 80GB)
    63→- Single request with KV-cache: 9x real-time
    64→- Batched (batch=64): 359x real-time throughput
    65→
    66→### Voice Prompting
    67→The inference engine supports audio prompts — a reference clip is encoded and used to seed autoregressive generation, effectively providing voice cloning without explicit speaker conditioning.
    68→
    69→## Training History
    70→
    71→### v1 — 80 mels + Griffin-Lim + Vocos
    72→- Simple MelDecoder (2x transposed conv), 80-mel output
    73→- Vocoder: Griffin-Lim -> Vocos refinement pipeline
    74→- Result: teacher-forced quality decent but "rustic" sound from GL artifacts
    75→
    76→### v2 — DAC-based
    77→- Experimental DAC codec variant
    78→- Checkpoints preserved at `v2/output/checkpoints/`
    79→
    80→### v3 — 100 mels + upgraded decoder + spectral loss (current)
    81→- Direct 100-mel output matching Vocos natively (no GL middleman)
    82→- Upgraded MelDecoder with ResConvBlock residual layers, 512 hidden channels
    83→- Multi-resolution spectral loss for sharp mel output
    84→- 24M parameters, training on `processed_data_100mel/`
    85→- Output: `output_v3/`
    86→
    87→## Usage
    88→
    89→### Preprocessing
    90→```bash
    91→python3 preprocess.py --input /path/to/dataset --output /path/to/processed_data
    92→```
    93→
    94→### Training
    95→```bash
    96→python3 train.py \
    97→  --data_dir processed_data_100mel \
    98→  --output_dir output_v3 \
    99→  --epochs 200 \
   100→  --batch_size 256 \
   101→  --spectral_weight 0.5
   102→```
   103→
   104→### Inference
   105→```bash
   106→# Synthesize from text
   107→python3 inference.py --checkpoint output_v3/checkpoints/best.pt --text "your text" --mode synthesize
   108→
   109→# Teacher-forced reconstruction
   110→python3 inference.py --checkpoint output_v3/checkpoints/best.pt --mel path/to/mel.pt --mode predict
   111→```
   112→
   113→### Batched Inference
   114→```bash
   115→python3 inference_engine.py --checkpoint output_v3/checkpoints/best.pt --benchmark
   116→```
   117→
   118→## Reference
   119→- LeWM paper: https://le-wm.github.io/
   120→

# LeWM TTS

JEPA-based Text-to-Speech system adapted from [LeWorldModel](https://le-wm.github.io/).

## Architecture

A 24M parameter model with 4 core components:

```
Text (Hindi/UTF-8 bytes) --> TextEncoder --> text embeddings ---+
                                                                |--> JEPAPredictor --> predicted embeddings
Audio (mel spectrogram)  --> AudioEncoder --> audio embeddings --+           |
                                                                            v
                                                                       MelDecoder --> mel spectrogram --> Vocos --> waveform
```

| Component | Description |
|-----------|-------------|
| **TextEncoder** | Byte-level character transformer (4 layers, 256d, 4 heads). Encodes Hindi text via UTF-8 bytes — no tokenizer needed. |
| **AudioEncoder** | 1D CNN (stride-4 downsample) + Transformer (4 layers). Outputs mu/logvar for Gaussian latent space with reparameterization. |
| **JEPAPredictor** | Transformer decoder (6 layers) with causal self-attention + cross-attention to text. Predicts next audio embedding autoregressively. Has KV-cache for fast inference. |
| **MelDecoder** | Linear projection to 512d, 2x transposed convolution (4x upsample), residual conv blocks (ResConvBlock), and refinement layers. Outputs 100-mel spectrogram. |
| **Vocoder** | [Vocos](https://github.com/gemelo-ai/vocos) `charactr/vocos-mel-24khz` — directly decodes 100-mel to 24kHz waveform. No Griffin-Lim. |

### Speaker Embedding (planned)
Built-in support for multi-speaker via learned embedding table added to text embeddings. Not yet trained — target is 10-20 speaker multi-speaker TTS.

## Training Losses

4 losses, all active during training:

1. **Prediction Loss (MSE)** — next-embedding prediction: `MSE(predicted[t], target[t+1])`
2. **KL Loss** — Gaussian regularizer: `KL(q(z|x) || N(0, I))` with annealing over first 20% of training
3. **Reconstruction Loss (L1)** — mel decoder output vs ground truth mel spectrogram
4. **Multi-Resolution Spectral Loss** — STFT at 3 resolutions (64/128/256 FFT sizes), combines spectral convergence + log magnitude L1. Forces sharp spectral detail instead of blurry averages.

## Dataset

- **Source**: 20 hours of studio-quality single-speaker Hindi audio
- **Preprocessing**: silence-based segmentation into 3-12s chunks, mel extraction
- **Processed**: 7,009 segments, 17.88 hours
- **Mel config**: 100 mels, n_fft=1024, hop=256, sr=24000, f_max=None (full bandwidth), log-clamped at 1e-7 — matches Vocos `mel-24khz` exactly

## Files

| File | Purpose |
|------|---------|
| `model.py` | All model components: TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, MultiResolutionSpectralLoss, LeWMTTS |
| `train.py` | Training loop with KL annealing, gradient clipping, TensorBoard logging, checkpointing |
| `inference.py` | Single-request inference — AR synthesis, teacher-forced reconstruction |
| `inference_engine.py` | Production batched inference engine — KV-cache, FP16, voice prompts, benchmarking |
| `preprocess.py` | Data preprocessing — audio segmentation, mel extraction, manifest creation |
| `dataset.py` | PyTorch dataset + dataloader with dynamic batching and padding |

## Inference

### Modes
- **Synthesize**: Pure autoregressive generation from text. Requires audio prompt (voice cloning) for best results.
- **Reconstruct**: Encode real mel -> decode back (tests encoder/decoder quality)
- **Predict (teacher-forced)**: Encode mel -> predictor -> decode (tests full pipeline without AR drift)

### Performance (A100 80GB)
- Single request with KV-cache: 9x real-time
- Batched (batch=64): 359x real-time throughput

### Voice Prompting
The inference engine supports audio prompts — a reference clip is encoded and used to seed autoregressive generation, effectively providing voice cloning without explicit speaker conditioning.

## Training History

### v1 — 80 mels + Griffin-Lim + Vocos
- Simple MelDecoder (2x transposed conv), 80-mel output
- Vocoder: Griffin-Lim -> Vocos refinement pipeline
- Result: teacher-forced quality decent but "rustic" sound from GL artifacts

### v2 — DAC-based
- Experimental DAC codec variant
- Checkpoints preserved at `v2/output/checkpoints/`

### v3 — 100 mels + upgraded decoder + spectral loss (current)
- Direct 100-mel output matching Vocos natively (no GL middleman)
- Upgraded MelDecoder with ResConvBlock residual layers, 512 hidden channels
- Multi-resolution spectral loss for sharp mel output
- 24M parameters, training on `processed_data_100mel/`
- Output: `output_v3/`

## Usage

### Preprocessing
```bash
python3 preprocess.py --input /path/to/dataset --output /path/to/processed_data
```

### Training
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v3 \
  --epochs 200 \
  --batch_size 256 \
  --spectral_weight 0.5
```

### Inference
```bash
# Synthesize from text
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --text "your text" --mode synthesize

# Teacher-forced reconstruction
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --mel path/to/mel.pt --mode predict
```

### Batched Inference
```bash
python3 inference_engine.py --checkpoint output_v3/checkpoints/best.pt --benchmark
```

## Reference
- LeWM paper: https://le-wm.github.io/


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

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
    }
    model, config = build_model(config)
    model = model.to(device)

    # ─── Build dataloader ───
    manifest_path = os.path.join(args.data_dir, "manifest.json")
    loader, dataset = build_dataloader(
        manifest_path, batch_size=args.batch_size,
        num_workers=args.num_workers,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    # ─── Optimizer & Scheduler ───
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.98), weight_decay=0.01,
    )

    # Warmup + cosine decay
    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

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

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

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

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

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

    # ─── Training loop ───
    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        epoch_loss = 0.0
        epoch_pred_loss = 0.0
        epoch_kl_loss = 0.0
        epoch_recon_loss = 0.0
        epoch_spectral_loss = 0.0
        t0 = time.time()

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

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

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_loss"]

            # Backward
            optimizer.zero_grad()
            total_loss.backward()

            # Gradient clipping
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            epoch_spectral_loss += spectral_loss.item()
            global_step += 1

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

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

        # Epoch summary
        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred_loss / n_batches
        avg_kl = epoch_kl_loss / n_batches
        avg_recon = epoch_recon_loss / n_batches
        avg_spectral = epoch_spectral_loss / n_batches
        epoch_time = time.time() - t0

        print(
            f"Epoch {epoch+1}/{args.epochs} | "
            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} kl={avg_kl:.4f} "
            f"recon={avg_recon:.4f} spec={avg_spectral:.4f} | {epoch_time:.1f}s"
        )

        # Save checkpoint
        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
            torch.save({
                "epoch": epoch,
                "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config,
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved checkpoint: {ckpt_path}")

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Data
    parser.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data")
    parser.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output")
    # Model
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--audio_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)
    parser.add_argument("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    parser.add_argument("--spectral_weight", type=float, default=0.5)
    # Training
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    # Logging
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=50)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args()
    train(args)


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

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

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

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

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

        self.max_mel_frames = max_mel_frames
        self.max_text_len = max_text_len

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

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

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

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

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

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

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

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

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

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

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

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


     1→"""
     2→LeWM TTS Model — JEPA-based Text-to-Speech
     3→
     4→Architecture:
     5→  - TextEncoder: Character-level transformer for Hindi/Devanagari text
     6→  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
     7→  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
     8→  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
     9→"""
    10→
    11→import math
    12→import torch
    13→import torch.nn as nn
    14→import torch.nn.functional as F
    15→import torchaudio
    16→
    17→
    18→# ─── Text Encoder ───────────────────────────────────────────────────────────
    19→
    20→class TextEncoder(nn.Module):
    21→    """Character-level transformer encoder for Hindi text."""
    22→
    23→    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
    24→        super().__init__()
    25→        self.d_model = d_model
    26→
    27→        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
    28→        self.char_embed = nn.Embedding(vocab_size, d_model)
    29→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
    30→        self.dropout = nn.Dropout(dropout)
    31→
    32→        encoder_layer = nn.TransformerEncoderLayer(
    33→            d_model=d_model, nhead=nhead,
    34→            dim_feedforward=d_model * 4,
    35→            dropout=dropout, batch_first=True,
    36→            activation="gelu",
    37→        )
    38→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    39→        self.proj = nn.Linear(d_model, d_model)
    40→
    41→    def forward(self, text_tokens, text_mask=None):
    42→        """
    43→        Args:
    44→            text_tokens: [B, T_text] long tensor of byte-level token IDs
    45→            text_mask: [B, T_text] bool, True = padding
    46→        Returns:
    47→            text_emb: [B, T_text, d_model]
    48→        """
    49→        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
    50→        x = self.pos_embed(x)
    51→        x = self.dropout(x)
    52→        x = self.transformer(x, src_key_padding_mask=text_mask)
    53→        x = self.proj(x)
    54→        return x
    55→
    56→
    57→# ─── Audio Encoder ──────────────────────────────────────────────────────────
    58→
    59→class AudioEncoder(nn.Module):
    60→    """
    61→    Encodes mel spectrograms into a sequence of audio embeddings.
    62→    Uses 1D CNN for downsampling + Transformer for contextualization.
    63→    """
    64→
    65→    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
    66→                 downsample_factor=4, dropout=0.1):
    67→        super().__init__()
    68→        self.d_model = d_model
    69→        self.downsample_factor = downsample_factor
    70→
    71→        # CNN downsampling: reduces temporal resolution by downsample_factor
    72→        if downsample_factor == 2:
    73→            self.conv_pre = nn.Sequential(
    74→                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    75→                nn.GELU(),
    76→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    77→                nn.GELU(),
    78→            )  # Total stride = 2
    79→        else:
    80→            self.conv_pre = nn.Sequential(
    81→                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    82→                nn.GELU(),
    83→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    84→                nn.GELU(),
    85→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    86→                nn.GELU(),
    87→            )  # Total stride = 4
    88→
    89→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
    90→        self.dropout = nn.Dropout(dropout)
    91→
    92→        encoder_layer = nn.TransformerEncoderLayer(
    93→            d_model=d_model, nhead=nhead,
    94→            dim_feedforward=d_model * 4,
    95→            dropout=dropout, batch_first=True,
    96→            activation="gelu",
    97→        )
    98→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    99→
   100→        # Project to embedding with mu and logvar for Gaussian regularizer
   101→        self.proj_mu = nn.Linear(d_model, d_model)
   102→        self.proj_logvar = nn.Linear(d_model, d_model)
   103→
   104→    def forward(self, mel, mel_mask=None):
   105→        """
   106→        Args:
   107→            mel: [B, n_mels, T_mel] — log mel spectrogram
   108→            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
   109→        Returns:
   110→            z: [B, T_down, d_model] — sampled latent embeddings
   111→            mu: [B, T_down, d_model]
   112→            logvar: [B, T_down, d_model]
   113→        """
   114→        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
   115→        x = self.conv_pre(mel)
   116→        x = x.transpose(1, 2)  # [B, T_down, d_model]
   117→
   118→        # Downsample mask if provided
   119→        if mel_mask is not None:
   120→            # Approximate downsampled mask
   121→            T_down = x.shape[1]
   122→            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
   123→
   124→        x = self.pos_embed(x)
   125→        x = self.dropout(x)
   126→        x = self.transformer(x, src_key_padding_mask=mel_mask)
   127→
   128→        mu = self.proj_mu(x)
   129→        logvar = self.proj_logvar(x)
   130→
   131→        # Reparameterization trick
   132→        if self.training:
   133→            std = torch.exp(0.5 * logvar)
   134→            eps = torch.randn_like(std)
   135→            z = mu + eps * std
   136→        else:
   137→            z = mu
   138→
   139→        return z, mu, logvar
   140→
   141→
   142→# ─── JEPA Predictor ─────────────────────────────────────────────────────────
   143→
   144→class JEPAPredictor(nn.Module):
   145→    """
   146→    Predicts next audio embedding given:
   147→    - Text context (from TextEncoder)
   148→    - Previous audio embeddings (from AudioEncoder)
   149→
   150→    Uses cross-attention to condition on text, and causal self-attention
   151→    over the audio embedding sequence.
   152→    """
   153→
   154→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
   155→        super().__init__()
   156→        self.d_model = d_model
   157→
   158→        self.audio_input_proj = nn.Linear(d_model, d_model)
   159→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
   160→        self.dropout = nn.Dropout(dropout)
   161→
   162→        # Transformer decoder with causal self-attention + cross-attention to text
   163→        decoder_layer = nn.TransformerDecoderLayer(
   164→            d_model=d_model, nhead=nhead,
   165→            dim_feedforward=d_model * 4,
   166→            dropout=dropout, batch_first=True,
   167→            activation="gelu",
   168→        )
   169→        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
   170→
   171→        # Output projection to predict next embedding
   172→        self.output_proj = nn.Sequential(
   173→            nn.Linear(d_model, d_model),
   174→            nn.GELU(),
   175→            nn.Linear(d_model, d_model),
   176→        )
   177→
   178→    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
   179→        """
   180→        Args:
   181→            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
   182→            text_emb: [B, T_text, d_model] — text embeddings
   183→            audio_mask: [B, T_audio] bool, True = padding
   184→            text_mask: [B, T_text] bool, True = padding
   185→        Returns:
   186→            predicted: [B, T_audio, d_model] — predicted next embeddings
   187→        """
   188→        x = self.audio_input_proj(audio_emb)
   189→        x = self.pos_embed(x)
   190→        x = self.dropout(x)
   191→
   192→        # Causal mask: each position can only attend to previous positions
   193→        T = x.shape[1]
   194→        causal_mask = torch.triu(
   195→            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
   196→        )
   197→
   198→        predicted = self.transformer(
   199→            x, text_emb,
   200→            tgt_mask=causal_mask,
   201→            tgt_key_padding_mask=audio_mask,
   202→            memory_key_padding_mask=text_mask,
   203→        )
   204→
   205→        predicted = self.output_proj(predicted)
   206→        return predicted
   207→
   208→    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
   209→        """Run multi-head attention with KV cache.
   210→        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
   211→        cache_k, cache_v: [B, T_prev, d] projected keys/values
   212→        Returns: output [B, 1, d], updated cache_k, updated cache_v
   213→        """
   214→        d = mha.embed_dim
   215→        nhead = mha.num_heads
   216→        head_dim = d // nhead
   217→        B = q_input.shape[0]
   218→
   219→        # Project Q from q_input, K/V from kv_input
   220→        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
   221→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   222→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   223→
   224→        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
   225→        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
   226→        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]
   227→
   228→        # Append to cache
   229→        if cache_k is not None and cache_k.shape[1] > 0:
   230→            k = torch.cat([cache_k, k_new], dim=1)
   231→            v = torch.cat([cache_v, v_new], dim=1)
   232→        else:
   233→            k = k_new
   234→            v = v_new
   235→
   236→        # Multi-head attention
   237→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
   238→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
   239→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   240→
   241→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   242→        attn = F.softmax(attn, dim=-1)
   243→        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]
   244→
   245→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   246→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   247→
   248→        return out, k, v
   249→
   250→    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
   251→        """Cross-attention with cached memory K,V (computed once).
   252→        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
   253→        d = mha.embed_dim
   254→        nhead = mha.num_heads
   255→        head_dim = d // nhead
   256→        B = q_input.shape[0]
   257→
   258→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   259→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   260→
   261→        q = F.linear(q_input, Wq, bq)  # [B, 1, d]
   262→
   263→        if cache_k is None:
   264→            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
   265→            v = F.linear(memory, Wv, bv)
   266→        else:
   267→            k, v = cache_k, cache_v
   268→
   269→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
   270→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
   271→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   272→
   273→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   274→        if memory_mask is not None:
   275→            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
   276→            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
   277→        attn = F.softmax(attn, dim=-1)
   278→        out = torch.matmul(attn, v_mh)
   279→
   280→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   281→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   282→
   283→        return out, k, v
   284→
   285→    def init_cache(self, num_layers):
   286→        """Initialize empty KV cache."""
   287→        return {
   288→            'self_k': [None] * num_layers,
   289→            'self_v': [None] * num_layers,
   290→            'cross_k': [None] * num_layers,
   291→            'cross_v': [None] * num_layers,
   292→        }
   293→
   294→    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   295→        """Process one token with KV cache.
   296→        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
   297→        text_mask: [B, T_text] bool, True = padding
   298→        Returns: predicted [B, 1, d_model], updated cache
   299→        """
   300→        x = self.audio_input_proj(new_emb)
   301→        # Add positional encoding for this specific position
   302→        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
   303→
   304→        for i, layer in enumerate(self.transformer.layers):
   305→            # Self-attention with cache (post-norm)
   306→            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
   307→                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
   308→            )
   309→            x = layer.norm1(x + layer.dropout1(sa_out))
   310→
   311→            # Cross-attention with cached text KV
   312→            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
   313→                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
   314→                memory_mask=text_mask,
   315→            )
   316→            x = layer.norm2(x + layer.dropout2(ca_out))
   317→
   318→            # FFN
   319→            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
   320→            x = layer.norm3(x + layer.dropout3(ff_out))
   321→
   322→        return self.output_proj(x), cache
   323→
   324→
   325→# ─── Positional Encoding ────────────────────────────────────────────────────
   326→
   327→class SinusoidalPositionalEncoding(nn.Module):
   328→    def __init__(self, d_model, max_len=8192):
   329→        super().__init__()
   330→        pe = torch.zeros(max_len, d_model)
   331→        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
   332→        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
   333→        pe[:, 0::2] = torch.sin(position * div_term)
   334→        pe[:, 1::2] = torch.cos(position * div_term)
   335→        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
   336→
   337→    def forward(self, x):
   338→        return x + self.pe[:, :x.shape[1]]
   339→
   340→
   341→# ─── Mel Decoder ────────────────────────────────────────────────────────────
   342→
   343→class ResConvBlock(nn.Module):
   344→    """Residual conv block for mel decoder."""
   345→
   346→    def __init__(self, channels, kernel_size=3):
   347→        super().__init__()
   348→        self.net = nn.Sequential(
   349→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   350→            nn.GELU(),
   351→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   352→        )
   353→
   354→    def forward(self, x):
   355→        return x + self.net(x)
   356→
   357→
   358→class MelDecoder(nn.Module):
   359→    """Decodes latent embeddings back to mel spectrograms.
   360→    Improved architecture with residual blocks for sharper mel output."""
   361→
   362→    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
   363→        super().__init__()
   364→        hidden = d_model * 2  # 512 — more capacity
   365→
   366→        self.proj = nn.Sequential(
   367→            nn.Linear(d_model, hidden),
   368→            nn.GELU(),
   369→        )
   370→        if upsample_factor == 2:
   371→            self.upsample = nn.Sequential(
   372→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   373→                nn.GELU(),
   374→                ResConvBlock(hidden, kernel_size=3),
   375→            )
   376→        else:
   377→            self.upsample = nn.Sequential(
   378→                # 4x upsample (2x + 2x)
   379→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   380→                nn.GELU(),
   381→                ResConvBlock(hidden, kernel_size=3),
   382→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   383→                nn.GELU(),
   384→                ResConvBlock(hidden, kernel_size=3),
   385→            )
   386→        # Refinement + projection to mel bins
   387→        self.out = nn.Sequential(
   388→            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
   389→            nn.GELU(),
   390→            ResConvBlock(hidden, kernel_size=5),
   391→            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
   392→        )
   393→
   394→    def forward(self, z):
   395→        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
   396→        x = self.proj(z)
   397→        x = x.transpose(1, 2)  # [B, hidden, T_down]
   398→        x = self.upsample(x)   # [B, hidden, T_up]
   399→        return self.out(x)      # [B, n_mels, T_up]
   400→
   401→
   402→# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────
   403→
   404→class MultiResolutionSpectralLoss(nn.Module):
   405→    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
   406→    This forces the mel decoder to produce sharp, well-defined spectral structure
   407→    rather than blurry averages that Vocos can't decode properly."""
   408→
   409→    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
   410→        super().__init__()
   411→        # Each resolution is (n_fft, hop_length, win_length)
   412→        self.resolutions = resolutions
   413→
   414→    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
   415→        """Compute spectral convergence + log magnitude loss for one resolution."""
   416→        window = torch.hann_window(win_length, device=pred.device)
   417→
   418→        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
   419→        # Flatten mel bins into batch: [B*n_mels, T]
   420→        B, M, T = pred.shape
   421→        pred_flat = pred.reshape(B * M, T)
   422→        target_flat = target.reshape(B * M, T)
   423→
   424→        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
   425→                               window=window, return_complex=True)
   426→        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
   427→                                 window=window, return_complex=True)
   428→
   429→        pred_mag = pred_stft.abs()
   430→        target_mag = target_stft.abs()
   431→
   432→        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
   433→        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)
   434→
   435→        # Log magnitude loss
   436→        log_mag_loss = F.l1_loss(
   437→            torch.log(pred_mag + 1e-7),
   438→            torch.log(target_mag + 1e-7),
   439→        )
   440→
   441→        return sc_loss + log_mag_loss
   442→
   443→    def forward(self, pred, target):
   444→        """pred, target: [B, n_mels, T]"""
   445→        loss = torch.tensor(0.0, device=pred.device)
   446→        count = 0
   447→        for n_fft, hop, win in self.resolutions:
   448→            # Skip if T is too short for this resolution
   449→            if pred.shape[2] >= n_fft:
   450→                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
   451→                count += 1
   452→        if count > 0:
   453→            loss = loss / count
   454→        return loss
   455→
   456→
   457→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   458→
   459→class LeWMTTS(nn.Module):
   460→    """
   461→    Complete LeWM TTS model.
   462→
   463→    Training: text + mel → audio embeddings → predict next embedding
   464→    Two losses:
   465→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   466→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   467→    """
   468→
   469→    def __init__(self, config):
   470→        super().__init__()
   471→        d_model = config.get("d_model", 256)
   472→        nhead = config.get("nhead", 4)
   473→        n_mels = config.get("n_mels", 100)
   474→        text_vocab = config.get("text_vocab_size", 256)
   475→        text_layers = config.get("text_encoder_layers", 4)
   476→        audio_layers = config.get("audio_encoder_layers", 4)
   477→        predictor_layers = config.get("predictor_layers", 6)
   478→        dropout = config.get("dropout", 0.1)
   479→        self.kl_weight = config.get("kl_weight", 0.01)
   480→        n_speakers = config.get("n_speakers", 1)
   481→
   482→        self.text_encoder = TextEncoder(
   483→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   484→            num_layers=text_layers, dropout=dropout,
   485→        )
   486→        downsample_factor = config.get("downsample_factor", 4)
   487→        self.audio_encoder = AudioEncoder(
   488→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   489→            num_layers=audio_layers, downsample_factor=downsample_factor,
   490→            dropout=dropout,
   491→        )
   492→        self.predictor = JEPAPredictor(
   493→            d_model=d_model, nhead=nhead,
   494→            num_layers=predictor_layers, dropout=dropout,
   495→        )
   496→
   497→        # Mel decoder: reconstructs mel from latent embeddings
   498→        # Needed so the latent space is grounded to actual audio
   499→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
   500→        self.recon_weight = config.get("recon_weight", 1.0)
   501→
   502→        # Speaker embedding (for multi-speaker)
   503→        self.n_speakers = n_speakers
   504→        if n_speakers > 1:
   505→            self.speaker_embed = nn.Embedding(n_speakers, d_model)
   506→        else:
   507→            self.speaker_embed = None
   508→
   509→        # Multi-resolution spectral loss for sharper mel output
   510→        self.spectral_loss = MultiResolutionSpectralLoss()
   511→        self.spectral_weight = config.get("spectral_weight", 0.5)
   512→
   513→        self.config = config
   514→
   515→    def _add_speaker(self, text_emb, speaker_id):
   516→        """Add speaker embedding to text embeddings."""
   517→        if self.speaker_embed is not None and speaker_id is not None:
   518→            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
   519→            text_emb = text_emb + spk
   520→        return text_emb
   521→
   522→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
   523→        """
   524→        Args:
   525→            mel: [B, n_mels, T_mel]
   526→            text_tokens: [B, T_text]
   527→            mel_mask: [B, T_mel] bool
   528→            text_mask: [B, T_text] bool
   529→        Returns:
   530→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   531→        """
   532→        # Encode text + speaker conditioning
   533→        text_emb = self.text_encoder(text_tokens, text_mask)
   534→        text_emb = self._add_speaker(text_emb, speaker_id)
   535→
   536→        # Encode audio → latent embeddings
   537→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   538→
   539→        # Predict next embedding: input z[:-1], target z[1:]
   540→        input_emb = z[:, :-1]     # [B, T-1, d]
   541→        target_emb = z[:, 1:]     # [B, T-1, d]
   542→
   543→        # Adjust mask for downsampled sequence
   544→        if mel_mask is not None:
   545→            T_down = z.shape[1]
   546→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   547→            pred_mask = ds_mask[:, :-1]
   548→        else:
   549→            pred_mask = None
   550→
   551→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   552→
   553→        # ─── Loss 1: Next-embedding prediction (MSE) ───
   554→        # Mask out padding positions
   555→        if pred_mask is not None:
   556→            valid = ~pred_mask  # True = valid
   557→            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
   558→            prediction_loss = F.mse_loss(
   559→                predicted * valid, target_emb * valid, reduction="sum"
   560→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   561→        else:
   562→            prediction_loss = F.mse_loss(predicted, target_emb)
   563→
   564→        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
   565→        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
   566→        if mel_mask is not None:
   567→            T_down = mu.shape[1]
   568→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   569→            valid = (~ds_mask).unsqueeze(-1).float()
   570→            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
   571→            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
   572→        else:
   573→            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
   574→
   575→        # ─── Loss 3: Mel reconstruction ───
   576→        # Decode latents back to mel and compare with input
   577→        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
   578→        # Trim/pad to match input mel length
   579→        T_mel = mel.shape[2]
   580→        T_recon = mel_recon.shape[2]
   581→        if T_recon > T_mel:
   582→            mel_recon = mel_recon[:, :, :T_mel]
   583→        elif T_recon < T_mel:
   584→            mel = mel[:, :, :T_recon]
   585→            if mel_mask is not None:
   586→                mel_mask = mel_mask[:, :T_recon]
   587→
   588→        if mel_mask is not None:
   589→            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
   590→            recon_loss = F.l1_loss(
   591→                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
   592→            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
   593→        else:
   594→            recon_loss = F.l1_loss(mel_recon, mel)
   595→
   596→        # ─── Loss 4: Multi-resolution spectral loss ───
   597→        # Apply on valid (non-padded) regions only
   598→        if mel_mask is not None:
   599→            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
   600→        else:
   601→            spectral_loss = self.spectral_loss(mel_recon, mel)
   602→
   603→        total_loss = (prediction_loss
   604→                      + self.kl_weight * kl_loss
   605→                      + self.recon_weight * recon_loss
   606→                      + self.spectral_weight * spectral_loss)
   607→
   608→        return {
   609→            "total_loss": total_loss,
   610→            "prediction_loss": prediction_loss,
   611→            "kl_loss": kl_loss,
   612→            "recon_loss": recon_loss,
   613→            "spectral_loss": spectral_loss,
   614→        }
   615→
   616→    def encode_audio(self, mel):
   617→        """Encode mel to latent embeddings (for inference)."""
   618→        z, mu, _ = self.audio_encoder(mel)
   619→        return mu  # Use mean at inference
   620→
   621→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   622→        """Predict next audio embedding autoregressively."""
   623→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   624→        return predicted[:, -1:]  # Only the last prediction
   625→
   626→    def init_ar_cache(self):
   627→        """Initialize KV cache for autoregressive inference."""
   628→        num_layers = len(self.predictor.transformer.layers)
   629→        return self.predictor.init_cache(num_layers)
   630→
   631→    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   632→        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
   633→        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
   634→
   635→
   636→def count_parameters(model):
   637→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   638→
   639→
   640→def build_model(config=None):
   641→    if config is None:
   642→        config = {
   643→            "d_model": 256,
   644→            "nhead": 4,
   645→            "n_mels": 100,
   646→            "text_vocab_size": 256,
   647→            "text_encoder_layers": 4,
   648→            "audio_encoder_layers": 4,
   649→            "predictor_layers": 6,
   650→            "dropout": 0.1,
   651→            "kl_weight": 0.01,
   652→        }
   653→    model = LeWMTTS(config)
   654→    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
   655→    return model, config
   656→

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

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

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

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        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):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        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)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        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)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

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

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    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)

        # Transformer decoder with causal self-attention + cross-attention to text
        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)

        # Output projection to predict next embedding
        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):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        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,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        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)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        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 = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        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)  # [B, nhead, 1, head_dim]

        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):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        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)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            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:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            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):
        """Initialize empty KV cache."""
        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):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            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))

            # Cross-attention with cached text KV
            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))

            # FFN
            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))  # [1, max_len, d_model]

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

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    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 MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

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

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

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

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

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

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

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

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

        self.config = config

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

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

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

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

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

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

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

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

        # ─── Loss 4: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

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

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

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

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

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

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

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

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


Now let me look at the training logs and the v3c version, as well as the inference code and any v2 files to understand what happened.


Now let me check the training logs, inference code, and the v3c output to understand what went wrong.


     1→#!/bin/bash
     2→# Wait for mels to finish regenerating, then resume v1 training
     3→
     4→echo "Waiting for mels to finish..."
     5→while [ $(ls /home/ubuntu/lewm-tts/processed_data/mels/ 2>/dev/null | wc -l) -lt 7009 ]; do
     6→    sleep 10
     7→done
     8→echo "Mels ready: $(ls /home/ubuntu/lewm-tts/processed_data/mels/ | wc -l)"
     9→
    10→echo "Starting v1 training from epoch 70, 1000 total epochs..."
    11→cd /home/ubuntu/lewm-tts
    12→python3 -u train.py \
    13→    --epochs 300 \
    14→    --batch_size 64 \
    15→    --lr 1e-4 \
    16→    --kl_weight 0.1 \
    17→    --recon_weight 1.0 \
    18→    --num_workers 4 \
    19→    --print_every 20 \
    20→    --log_every 5 \
    21→    --save_every 50 \
    22→    --resume output/checkpoints/epoch_0070.pt \
    23→    --output_dir /home/ubuntu/lewm-tts/output
    24→

#!/bin/bash
# Wait for mels to finish regenerating, then resume v1 training

echo "Waiting for mels to finish..."
while [ $(ls /home/ubuntu/lewm-tts/processed_data/mels/ 2>/dev/null | wc -l) -lt 7009 ]; do
    sleep 10
done
echo "Mels ready: $(ls /home/ubuntu/lewm-tts/processed_data/mels/ | wc -l)"

echo "Starting v1 training from epoch 70, 1000 total epochs..."
cd /home/ubuntu/lewm-tts
python3 -u train.py \
    --epochs 300 \
    --batch_size 64 \
    --lr 1e-4 \
    --kl_weight 0.1 \
    --recon_weight 1.0 \
    --num_workers 4 \
    --print_every 20 \
    --log_every 5 \
    --save_every 50 \
    --resume output/checkpoints/epoch_0070.pt \
    --output_dir /home/ubuntu/lewm-tts/output