status


The user is asking for a status update. Let me check my memory files for context.


ToolSearch

Question: select:Read


Question: select:Read


     1→# LeWM TTS Project
     2→
     3→## Overview
     4→Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
     5→- Paper: https://le-wm.github.io/
     6→- Project folder: `/home/ubuntu/lewm-tts`
     7→- Goal: 10-20 speaker multi-speaker TTS on this 14M→24M model
     8→
     9→## Architecture (v3 — Current)
    10→- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
    11→- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
    12→- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
    13→- **MelDecoder**: Upgraded — Linear→512 + 2x ConvTranspose1d + ResConvBlock residual blocks + refinement conv
    14→- **Vocoder**: Direct Vocos decode (model outputs 100 mels matching `charactr/vocos-mel-24khz`)
    15→- **4 losses**: prediction (MSE), KL (Gaussian reg), reconstruction (L1 mel), **multi-resolution spectral loss**
    16→- **Speaker embedding**: support for n_speakers > 1 (not yet trained)
    17→- **Model size**: 24.00M parameters
    18→- **KL annealing**: ramps up over first 20% of training
    19→
    20→## Multi-Resolution Spectral Loss
    21→- STFT at 3 resolutions: (64,16,64), (128,32,128), (256,64,256)
    22→- Applied per mel-bin (flattens [B, n_mels, T] → [B*n_mels, T])
    23→- Spectral convergence + log magnitude L1
    24→- Weight: 0.5 (configurable via --spectral_weight)
    25→
    26→## Files
    27→- `preprocess.py` — segments audio, extracts mels, saves manifest
    28→- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder (upgraded), MultiResolutionSpectralLoss, LeWMTTS
    29→- `dataset.py` — TTSDataset, collate_fn, build_dataloader
    30→- `train.py` — training loop with logging, checkpointing, KL annealing, spectral loss
    31→- `inference.py` — AR generation with Vocos vocoder (single request)
    32→- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts
    33→
    34→## Dataset
    35→- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
    36→- **100-mel processed**: `/home/ubuntu/lewm-tts/processed_data_100mel` (7,009 segments, 17.88h)
    37→- Mel config: n_fft=1024, hop=256, n_mels=100, sr=24000, f_max=None, power=1.0, clip_val=1e-7
    38→
    39→## Training History
    40→- **v1**: 80 mels, simple decoder, GL→Vocos pipeline. Rustic sound.
    41→- **v2**: DAC-based, checkpoints preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`
    42→- **v3** (current): 100 mels, upgraded MelDecoder + spectral loss, direct Vocos. Training at output_v3/.
    43→
    44→## Inference Findings
    45→- **Predict (teacher-forced)**: works well, good quality
    46→- **AR synthesis**: broken with random start token — needs audio prompt (voice cloning)
    47→- **Best early checkpoints**: epoch 70-100 (later epochs collapse for AR)
    48→- Noise tolerance test: GT mel + noise at L1=0.3 sounds natural; model's structured error sounds worse
    49→
    50→## Inference Performance (A100 80GB)
    51→- KV-cache: 9x real-time single request
    52→- Batch=64: 359x throughput
    53→
    54→## Next Steps
    55→- Wait for v3 training to converge, test quality
    56→- If mel quality good with Vocos, proceed to multi-speaker training
    57→- Scheduled sampling during training (robustify AR)
    58→

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`
- Goal: 10-20 speaker multi-speaker TTS on this 14M→24M model

## Architecture (v3 — Current)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
- **MelDecoder**: Upgraded — Linear→512 + 2x ConvTranspose1d + ResConvBlock residual blocks + refinement conv
- **Vocoder**: Direct Vocos decode (model outputs 100 mels matching `charactr/vocos-mel-24khz`)
- **4 losses**: prediction (MSE), KL (Gaussian reg), reconstruction (L1 mel), **multi-resolution spectral loss**
- **Speaker embedding**: support for n_speakers > 1 (not yet trained)
- **Model size**: 24.00M parameters
- **KL annealing**: ramps up over first 20% of training

## Multi-Resolution Spectral Loss
- STFT at 3 resolutions: (64,16,64), (128,32,128), (256,64,256)
- Applied per mel-bin (flattens [B, n_mels, T] → [B*n_mels, T])
- Spectral convergence + log magnitude L1
- Weight: 0.5 (configurable via --spectral_weight)

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder (upgraded), MultiResolutionSpectralLoss, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing, spectral loss
- `inference.py` — AR generation with Vocos vocoder (single request)
- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **100-mel processed**: `/home/ubuntu/lewm-tts/processed_data_100mel` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=100, sr=24000, f_max=None, power=1.0, clip_val=1e-7

## Training History
- **v1**: 80 mels, simple decoder, GL→Vocos pipeline. Rustic sound.
- **v2**: DAC-based, checkpoints preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`
- **v3** (current): 100 mels, upgraded MelDecoder + spectral loss, direct Vocos. Training at output_v3/.

## Inference Findings
- **Predict (teacher-forced)**: works well, good quality
- **AR synthesis**: broken with random start token — needs audio prompt (voice cloning)
- **Best early checkpoints**: epoch 70-100 (later epochs collapse for AR)
- Noise tolerance test: GT mel + noise at L1=0.3 sounds natural; model's structured error sounds worse

## Inference Performance (A100 80GB)
- KV-cache: 9x real-time single request
- Batch=64: 359x throughput

## Next Steps
- Wait for v3 training to converge, test quality
- If mel quality good with Vocos, proceed to multi-speaker training
- Scheduled sampling during training (robustify AR)


Here's the current status:

**LeWM TTS** — JEPA-based TTS, 24M params

- **Architecture (v3)**: Byte-level text encoder + Audio encoder (VAE) + JEPA predictor (6-layer transformer w/ KV-cache) + upgraded MelDecoder + Vocos vocoder
- **Dataset**: 7,009 segments, ~18h Hindi single-speaker audio
- **Training**: v3 in progress at `output_v3/` — upgraded decoder with multi-resolution spectral loss
- **Known issues**: AR synthesis needs audio prompt (voice cloning style); random start token doesn't work. Teacher-forced prediction works well.
- **Next steps**: Wait for v3 convergence, then test quality and move to multi-speaker training + scheduled sampling

Want me to check the current training progress or anything specific?


'/home/ubuntu/lewm-tts/sample_synth_v3.wav''/home/ubuntu/lewm-tts/sample_reconstruct_v3.wav' u ran infernce and we got this .. but these are noise only stil


The user is saying that the v3 inference produced only noise. Let me listen to the samples and investigate what's going on. Let me check the current state of training and the inference code to understand why.


Question: select:Bash,Glob


Question: select:Bash,Glob


Let me check the training status and look at the inference code to understand why the output is noise.


Let me investigate. I'll check the training state and inference code.


Bash

ls -la /home/ubuntu/lewm-tts/output_v3/checkpoints/ 2>/dev/null | tail -5

Check latest v3 checkpoints


ls -la /home/ubuntu/lewm-tts/sample_synth_v3.wav /home/ubuntu/lewm-tts/sample_reconstruct_v3.wav 2>/dev/null

Check the sample wav files


     1→"""
     2→Inference pipeline for LeWM TTS.
     3→Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
     4→"""
     5→
     6→import torch
     7→import torchaudio
     8→import numpy as np
     9→import argparse
    10→from pathlib import Path
    11→
    12→from model import LeWMTTS
    13→
    14→
    15→class LeWMTTSInference:
    16→    """Full inference pipeline."""
    17→
    18→    def __init__(self, checkpoint_path, device="cuda"):
    19→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    20→
    21→        # Load checkpoint
    22→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    23→        config = ckpt["config"]
    24→
    25→        # Build model
    26→        self.model = LeWMTTS(config).to(self.device)
    27→        self.model.load_state_dict(ckpt["model"])
    28→        self.model.eval()
    29→
    30→        # Vocos vocoder (directly decodes 100-mel)
    31→        from vocos import Vocos
    32→        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
    33→        self.vocos.eval()
    34→
    35→        self.config = config
    36→        self._start_emb_scale = 5.0
    37→        print(f"Model loaded from {checkpoint_path}")
    38→
    39→    def vocos_decode(self, mel):
    40→        """Decode mel [B, 100, T] → waveform via Vocos."""
    41→        with torch.no_grad():
    42→            waveform = self.vocos.decode(mel)  # [B, T_audio]
    43→        return waveform
    44→
    45→    def text_to_tokens(self, text):
    46→        tokens = list(text.encode("utf-8"))
    47→        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
    48→
    49→    @torch.no_grad()
    50→    def synthesize(self, text, max_steps=300, temperature=0.0):
    51→        """Generate speech from text using AR prediction with KV-cache."""
    52→        text_tokens = self.text_to_tokens(text)
    53→        text_emb = self.model.text_encoder(text_tokens)
    54→
    55→        d = self.config["d_model"]
    56→        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
    57→        all_embeddings = [start_emb]
    58→
    59→        cache = self.model.init_ar_cache()
    60→        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
    61→        if temperature > 0:
    62→            next_emb = next_emb + torch.randn_like(next_emb) * temperature
    63→        all_embeddings.append(next_emb)
    64→
    65→        prev_norms = [next_emb.norm().item()]
    66→
    67→        for step in range(1, max_steps):
    68→            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
    69→            if temperature > 0:
    70→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
    71→            all_embeddings.append(next_emb)
    72→
    73→            norm = next_emb.norm().item()
    74→            prev_norms.append(norm)
    75→            if len(prev_norms) > 20:
    76→                recent = prev_norms[-20:]
    77→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
    78→                    break
    79→
    80→        audio_embs = torch.cat(all_embeddings, dim=1)
    81→        mel = self.model.mel_decoder(audio_embs)
    82→        waveform = self.vocos_decode(mel.cpu())
    83→        return waveform.squeeze().numpy(), 24000
    84→
    85→    @torch.no_grad()
    86→    def reconstruct(self, mel_path):
    87→        """Teacher-forced reconstruction: encode real mel → decode back."""
    88→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
    89→
    90→        z, mu, logvar = self.model.audio_encoder(mel)
    91→        mel_recon = self.model.mel_decoder(mu)
    92→
    93→        waveform = self.vocos_decode(mel_recon.cpu())
    94→        return waveform.squeeze().numpy(), 24000
    95→
    96→    @torch.no_grad()
    97→    def reconstruct_with_prediction(self, mel_path):
    98→        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
    99→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
   100→
   101→        z, mu, logvar = self.model.audio_encoder(mel)
   102→
   103→        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
   104→        text_emb = self.model.text_encoder(text_tokens)
   105→
   106→        predicted = self.model.predictor(mu[:, :-1], text_emb)
   107→        combined = torch.cat([mu[:, :1], predicted], dim=1)
   108→
   109→        mel_recon = self.model.mel_decoder(combined)
   110→        waveform = self.vocos_decode(mel_recon.cpu())
   111→        return waveform.squeeze().numpy(), 24000
   112→
   113→    def save_audio(self, waveform, sr, output_path):
   114→        # Normalize
   115→        mx = np.abs(waveform).max()
   116→        if mx > 0:
   117→            waveform = waveform / mx * 0.95
   118→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
   119→        torchaudio.save(output_path, wav_tensor, sr)
   120→        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
   121→
   122→
   123→def main():
   124→    parser = argparse.ArgumentParser()
   125→    parser.add_argument("--checkpoint", required=True)
   126→    parser.add_argument("--text", default=None, help="Text to synthesize")
   127→    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
   128→    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
   129→                        default="synthesize")
   130→    parser.add_argument("--output", default="output.wav")
   131→    parser.add_argument("--max_steps", type=int, default=300)
   132→    parser.add_argument("--temperature", type=float, default=0.0)
   133→    args = parser.parse_args()
   134→
   135→    tts = LeWMTTSInference(args.checkpoint)
   136→
   137→    if args.mode == "synthesize":
   138→        assert args.text, "Need --text for synthesis"
   139→        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
   140→    elif args.mode == "reconstruct":
   141→        assert args.mel, "Need --mel for reconstruction"
   142→        waveform, sr = tts.reconstruct(args.mel)
   143→    elif args.mode == "predict":
   144→        assert args.mel, "Need --mel for prediction test"
   145→        waveform, sr = tts.reconstruct_with_prediction(args.mel)
   146→
   147→    tts.save_audio(waveform, sr, args.output)
   148→
   149→
   150→if __name__ == "__main__":
   151→    main()
   152→

"""
Inference pipeline for LeWM TTS.
Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
"""

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

from model import LeWMTTS

class LeWMTTSInference:
    """Full inference pipeline."""

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

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

        # Build model
        self.model = LeWMTTS(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # Vocos vocoder (directly decodes 100-mel)
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

        self.config = config
        self._start_emb_scale = 5.0
        print(f"Model loaded from {checkpoint_path}")

    def vocos_decode(self, mel):
        """Decode mel [B, 100, T] → waveform via Vocos."""
        with torch.no_grad():
            waveform = self.vocos.decode(mel)  # [B, T_audio]
        return waveform

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

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0):
        """Generate speech from text using AR prediction with KV-cache."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

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

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

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

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

        audio_embs = torch.cat(all_embeddings, dim=1)
        mel = self.model.mel_decoder(audio_embs)
        waveform = self.vocos_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """Teacher-forced reconstruction: encode real mel → decode back."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)
        mel_recon = self.model.mel_decoder(mu)

        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)

        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        predicted = self.model.predictor(mu[:, :-1], text_emb)
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        mel_recon = self.model.mel_decoder(combined)
        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

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

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

    tts = LeWMTTSInference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text for synthesis"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.mel, "Need --mel for reconstruction"
        waveform, sr = tts.reconstruct(args.mel)
    elif args.mode == "predict":
        assert args.mel, "Need --mel for prediction test"
        waveform, sr = tts.reconstruct_with_prediction(args.mel)

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

if __name__ == "__main__":
    main()


drwxrwxr-x 2 ubuntu ubuntu      4096 Mar 24 06:06 .
drwxrwxr-x 4 ubuntu ubuntu      4096 Mar 24 05:57 ..
-rw-rw-r-- 1 ubuntu ubuntu 106556117 Mar 24 06:10 best.pt
-rw-rw-r-- 1 ubuntu ubuntu 298760488 Mar 24 06:02 epoch_0010.pt
-rw-rw-r-- 1 ubuntu ubuntu 298760488 Mar 24 06:06 epoch_0020.pt


-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 06:07 /home/ubuntu/lewm-tts/sample_reconstruct_v3.wav
-rw-rw-r-- 1 ubuntu ubuntu 161358 Mar 24 06:07 /home/ubuntu/lewm-tts/sample_synth_v3.wav


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
ckpt = torch.load('output_v3/checkpoints/best.pt', map_location='cpu', weights_only=False)
print('Epoch:', ckpt.get('epoch', '?'))
print('Loss:', ckpt.get('best_loss', ckpt.get('loss', '?')))
print('Config keys:', list(ckpt['config'].keys()))
print('Config:', ckpt['config'])
"

Inspect best checkpoint metadata


Epoch: 28
Loss: 1.3300045993593004
Config keys: ['d_model', 'nhead', 'n_mels', 'text_vocab_size', 'text_encoder_layers', 'audio_encoder_layers', 'predictor_layers', 'dropout', 'kl_weight', 'recon_weight', 'spectral_weight']
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.1, 'recon_weight': 1.0, 'spectral_weight': 0.5}


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

"""
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
        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):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        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,
        )
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, 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)
        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


The model is only at epoch 28 with loss 1.33. That's very early in training. Let me check what the reconstruct function does - it's going through audio_encoder → mel_decoder, which is the VAE path. If even reconstruction is noise, the issue is likely:

1. The model hasn't trained enough (only 20 epochs saved, best at epoch 28)
2. Something wrong with how Vocos is being called

Let me check the actual mel output values and see what's happening.


python3 -c "
import torch
import numpy as np

# Load checkpoint
ckpt = torch.load('output_v3/checkpoints/best.pt', map_location='cpu', weights_only=False)
config = ckpt['config']

from model import LeWMTTS
model = LeWMTTS(config)
model.load_state_dict(ckpt['model'])
model.eval()

# Load a real mel for comparison
import glob
mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))[:1]
mel = torch.load(mel_files[0], weights_only=True)
print('Real mel shape:', mel.shape, 'min:', mel.min().item(), 'max:', mel.max().item(), 'mean:', mel.mean().item())

# Reconstruct through audio encoder → mel decoder
with torch.no_grad():
    mel_in = mel.unsqueeze(0)
    z, mu, logvar = model.audio_encoder(mel_in)
    print('Latent z shape:', z.shape, 'mean:', mu.mean().item(), 'std:', mu.std().item())

    mel_recon = model.mel_decoder(mu)
    print('Recon mel shape:', mel_recon.shape, 'min:', mel_recon.min().item(), 'max:', mel_recon.max().item(), 'mean:', mel_recon.mean().item())

    # L1 error
    T = min(mel_in.shape[2], mel_recon.shape[2])
    l1 = (mel_in[:,:,:T] - mel_recon[:,:,:T]).abs().mean().item()
    print('L1 recon error:', l1)
" 2>&1

Analyze mel reconstruction quality at current checkpoint


Real mel shape: torch.Size([100, 893]) min: -8.690321922302246 max: 4.463228225708008 mean: -2.1474547386169434
Latent z shape: torch.Size([1, 223, 256]) mean: -0.0008768412517383695 std: 0.28225278854370117
Recon mel shape: torch.Size([1, 100, 892]) min: -7.174448490142822 max: 3.0661306381225586 mean: -1.9635246992111206
L1 recon error: 0.7281765937805176


python3 -c "
import torch
import numpy as np
import glob

ckpt = torch.load('output_v3/checkpoints/best.pt', map_location='cpu', weights_only=False)
config = ckpt['config']

from model import LeWMTTS
model = LeWMTTS(config)
model.load_state_dict(ckpt['model'])
model.eval()

# Load vocos and test
from vocos import Vocos
vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz')
vocos.eval()

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))[:1]
mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0)

# 1. Direct vocos on real mel
with torch.no_grad():
    wav_real = vocos.decode(mel)
    print('Real mel -> vocos wav:', wav_real.shape, 'range:', wav_real.min().item(), wav_real.max().item())

    # 2. Reconstruct then vocos
    z, mu, logvar = model.audio_encoder(mel)
    mel_recon = model.mel_decoder(mu)
    wav_recon = vocos.decode(mel_recon)
    print('Recon mel -> vocos wav:', wav_recon.shape, 'range:', wav_recon.min().item(), wav_recon.max().item())

    # Check if vocos expects specific mel range
    print()
    print('Vocos expects log mel in what range?')
    print('Real mel range: [{:.2f}, {:.2f}]'.format(mel.min().item(), mel.max().item()))
    print('Recon mel range: [{:.2f}, {:.2f}]'.format(mel_recon.min().item(), mel_recon.max().item()))

    # Save both for comparison
    import torchaudio
    torchaudio.save('/tmp/test_real_vocos.wav', wav_real.cpu(), 24000)
    torchaudio.save('/tmp/test_recon_vocos.wav', wav_recon.cpu(), 24000)
    print('Saved test wavs to /tmp/')
" 2>&1

Test Vocos decoding on real vs reconstructed mel


Real mel -> vocos wav: torch.Size([1, 228352]) range: -0.8559535145759583 0.7635495066642761
Recon mel -> vocos wav: torch.Size([1, 228096]) range: -0.43805941939353943 0.36736735701560974

Vocos expects log mel in what range?
Real mel range: [-8.69, 4.46]
Recon mel range: [-7.17, 3.07]
Saved test wavs to /tmp/


python3 -c "
# Check how the inference script was called - look at the wav files
import torchaudio

for f in ['sample_synth_v3.wav', 'sample_reconstruct_v3.wav']:
    wav, sr = torchaudio.load(f'/home/ubuntu/lewm-tts/{f}')
    print(f'{f}: sr={sr}, shape={wav.shape}, range=[{wav.min():.4f}, {wav.max():.4f}], rms={wav.pow(2).mean().sqrt():.4f}')

# Also check the test files
for f in ['/tmp/test_real_vocos.wav', '/tmp/test_recon_vocos.wav']:
    wav, sr = torchaudio.load(f)
    print(f'{f}: sr={sr}, shape={wav.shape}, range=[{wav.min():.4f}, {wav.max():.4f}], rms={wav.pow(2).mean().sqrt():.4f}')
" 2>&1

Compare wav file stats between user's samples and test files


sample_synth_v3.wav: sr=24000, shape=torch.Size([1, 80640]), range=[-0.4084, 0.9500], rms=0.1053
sample_reconstruct_v3.wav: sr=24000, shape=torch.Size([1, 228096]), range=[-0.9379, 0.9500], rms=0.1141
/tmp/test_real_vocos.wav: sr=24000, shape=torch.Size([1, 228352]), range=[-0.8560, 0.7635], rms=0.0821
/tmp/test_recon_vocos.wav: sr=24000, shape=torch.Size([1, 228096]), range=[-0.4380, 0.3674], rms=0.0403


The reconstruction wav has reasonable RMS and range - it's not zero or extreme values. But the user says it sounds like noise. Let me check:

1. L1 recon error is 0.73 - that's quite high for a mel that ranges from -8.69 to 4.46 (range of ~13). So relative error is about 0.73/13 = 5.6%. That's not terrible but also not great.

2. The model has only trained 20 epochs (best at epoch 28 but only 10 and 20 saved). This is very early training.

3. Let me check if there's a mismatch between how mels are computed for training vs how Vocos expects them. Vocos expects a specific mel format.

4. Now I need to verify the mel preprocessing pipeline and ensure it matches what Vocos is designed to work with.


python3 -c "
# Check what Vocos actually expects vs what we're producing
from vocos import Vocos
import torch
vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz')

# Check vocos mel transform
print('Vocos feature extractor type:', type(vocos.feature_extractor))
print(dir(vocos.feature_extractor))

# Check if vocos has its own mel transform
import inspect
src = inspect.getsource(type(vocos.feature_extractor))
print(src[:2000])
" 2>&1

Check Vocos mel expectations


Vocos feature extractor type: <class 'vocos.feature_extractors.MelSpectrogramFeatures'>
['T_destination', '__annotations__', '__call__', '__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattr__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setstate__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_apply', '_backward_hooks', '_backward_pre_hooks', '_buffers', '_call_impl', '_compiled_call_impl', '_forward_hooks', '_forward_hooks_always_called', '_forward_hooks_with_kwargs', '_forward_pre_hooks', '_forward_pre_hooks_with_kwargs', '_get_backward_hooks', '_get_backward_pre_hooks', '_get_name', '_is_full_backward_hook', '_load_from_state_dict', '_load_state_dict_post_hooks', '_load_state_dict_pre_hooks', '_maybe_warn_non_full_backward_hook', '_modules', '_named_members', '_non_persistent_buffers_set', '_parameters', '_register_load_state_dict_pre_hook', '_register_state_dict_hook', '_replicate_for_data_parallel', '_save_to_state_dict', '_slow_forward', '_state_dict_hooks', '_state_dict_pre_hooks', '_version', '_wrapped_call_impl', 'add_module', 'apply', 'bfloat16', 'buffers', 'call_super_init', 'children', 'compile', 'cpu', 'cuda', 'double', 'dump_patches', 'eval', 'extra_repr', 'float', 'forward', 'get_buffer', 'get_extra_state', 'get_parameter', 'get_submodule', 'half', 'ipu', 'load_state_dict', 'mel_spec', 'modules', 'mtia', 'named_buffers', 'named_children', 'named_modules', 'named_parameters', 'padding', 'parameters', 'register_backward_hook', 'register_buffer', 'register_forward_hook', 'register_forward_pre_hook', 'register_full_backward_hook', 'register_full_backward_pre_hook', 'register_load_state_dict_post_hook', 'register_load_state_dict_pre_hook', 'register_module', 'register_parameter', 'register_state_dict_post_hook', 'register_state_dict_pre_hook', 'requires_grad_', 'set_extra_state', 'set_submodule', 'share_memory', 'state_dict', 'to', 'to_empty', 'train', 'training', 'type', 'xpu', 'zero_grad']
class MelSpectrogramFeatures(FeatureExtractor):
    def __init__(self, sample_rate=24000, n_fft=1024, hop_length=256, n_mels=100, padding="center"):
        super().__init__()
        if padding not in ["center", "same"]:
            raise ValueError("Padding must be 'center' or 'same'.")
        self.padding = padding
        self.mel_spec = torchaudio.transforms.MelSpectrogram(
            sample_rate=sample_rate,
            n_fft=n_fft,
            hop_length=hop_length,
            n_mels=n_mels,
            center=padding == "center",
            power=1,
        )

    def forward(self, audio, **kwargs):
        if self.padding == "same":
            pad = self.mel_spec.win_length - self.mel_spec.hop_length
            audio = torch.nn.functional.pad(audio, (pad // 2, pad // 2), mode="reflect")
        mel = self.mel_spec(audio)
        features = safe_log(mel)
        return features