1	"""
2	Dataset for codec-based JEPA TTS v5.
3	Loads precomputed EnCodec embeddings + text.
4	"""
5
6	import json
7	import torch
8	from torch.utils.data import Dataset, DataLoader
9
10
11	class CodecDataset(Dataset):
12	    def __init__(self, manifest_path, max_codec_frames=900):
13	        with open(manifest_path, "r") as f:
14	            self.manifest = json.load(f)
15	        self.max_codec_frames = max_codec_frames
16	        # Filter out too-long samples
17	        self.manifest = [e for e in self.manifest if e["emb_frames"] <= max_codec_frames]
18	        print(f"CodecDataset: {len(self.manifest)} samples (max {max_codec_frames} frames)")
19
20	    def __len__(self):
21	        return len(self.manifest)
22
23	    def __getitem__(self, idx):
24	        entry = self.manifest[idx]
25	        emb = torch.load(entry["emb_path"], weights_only=True)  # [128, T]
26	        text_tokens = torch.tensor(list(entry["text"].encode("utf-8")), dtype=torch.long)
27	        return {"codec_emb": emb, "text_tokens": text_tokens}
28
29
30	def collate_fn(batch):
31	    codec_embs = [b["codec_emb"] for b in batch]
32	    text_tokens = [b["text_tokens"] for b in batch]
33
34	    # Pad codec embeddings
35	    codec_dim = codec_embs[0].shape[0]
36	    max_codec_len = max(e.shape[1] for e in codec_embs)
37	    padded_codec = torch.zeros(len(batch), codec_dim, max_codec_len)
38	    codec_mask = torch.ones(len(batch), max_codec_len, dtype=torch.bool)
39	    for i, e in enumerate(codec_embs):
40	        T = e.shape[1]
41	        padded_codec[i, :, :T] = e
42	        codec_mask[i, :T] = False
43
44	    # Pad text
45	    max_text_len = max(t.shape[0] for t in text_tokens)
46	    padded_text = torch.zeros(len(batch), max_text_len, dtype=torch.long)
47	    text_mask = torch.ones(len(batch), max_text_len, dtype=torch.bool)
48	    for i, t in enumerate(text_tokens):
49	        padded_text[i, :t.shape[0]] = t
50	        text_mask[i, :t.shape[0]] = False
51
52	    return {
53	        "codec_emb": padded_codec,
54	        "text_tokens": padded_text,
55	        "codec_mask": codec_mask,
56	        "text_mask": text_mask,
57	    }
58
59
60	def build_dataloader(manifest_path, batch_size=64, num_workers=4, max_codec_frames=900):
61	    dataset = CodecDataset(manifest_path, max_codec_frames=max_codec_frames)
62	    loader = DataLoader(
63	        dataset, batch_size=batch_size, shuffle=True,
64	        num_workers=num_workers, collate_fn=collate_fn,
65	        pin_memory=True, drop_last=True,
66	    )
67	    return loader, dataset
68


Now let me check if there's a README that might contain additional context:


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


Now let me check the v2 directory structure and look for any additional config files:


Bash

find /home/ubuntu/lewm-tts -name "*.json" -o -name "*.yaml" -o -name "*.toml" | grep -v ".cache" | head -20


/home/ubuntu/lewm-tts/output_v3/config.json
/home/ubuntu/lewm-tts/output/config.json
/home/ubuntu/lewm-tts/processed_data_codec/manifest.json
/home/ubuntu/lewm-tts/output_v4/config.json
/home/ubuntu/lewm-tts/output_v5_pred_mask/config.json
/home/ubuntu/lewm-tts/output_v3c/config.json
/home/ubuntu/lewm-tts/v2/output/config.json
/home/ubuntu/lewm-tts/v2/processed_data/manifest.json
/home/ubuntu/lewm-tts/v2/processed_data/config.json
/home/ubuntu/lewm-tts/output_v5_pred_dur/config.json
/home/ubuntu/lewm-tts/output_v5/config.json
/home/ubuntu/lewm-tts/output_v5_pred/config.json
/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json
/home/ubuntu/lewm-tts/processed_data_100mel/config.json
/home/ubuntu/lewm-tts/processed_data/manifest.json
/home/ubuntu/lewm-tts/processed_data/config.json
/home/ubuntu/lewm-tts/output_v3b/config.json


Let me check a few config files to see the exact model dimensions:


1	{
2	  "d_model": 256,
3	  "nhead": 4,
4	  "n_mels": 100,
5	  "text_vocab_size": 256,
6	  "text_encoder_layers": 4,
7	  "audio_encoder_layers": 4,
8	  "predictor_layers": 6,
9	  "dropout": 0.1,
10	  "kl_weight": 0.1,
11	  "recon_weight": 1.0,
12	  "spectral_weight": 0.5
13	}


Let me check the v5 config:


1	{
2	  "d_model": 256,
3	  "nhead": 4,
4	  "codec_dim": 128,
5	  "text_vocab_size": 256,
6	  "text_encoder_layers": 4,
7	  "predictor_layers": 6,
8	  "dropout": 0.1,
9	  "pred_weight": 10.0,
10	  "roundtrip_weight": 1.0,
11	  "ema_decay": 0.998,
12	  "input_noise": 0.0
13	}


Let me check if there's a preprocess file and look at any logs:


1	"""
2	Data preprocessing for LeWM TTS.
3	- Loads Modi dataset from HuggingFace Arrow format
4	- Segments long clips into 3-15 second chunks using silence detection
5	- Extracts mel spectrograms
6	- Saves processed dataset
7	"""
8
9	import os
10	import json
11	import numpy as np
12	import torch
13	import torchaudio
14	from datasets import load_from_disk
15	from pathlib import Path
16	import argparse
17
18
19	# ─── Mel spectrogram config ───
20	MEL_CONFIG = {
21	    "sample_rate": 24000,
22	    "n_fft": 1024,
23	    "hop_length": 256,
24	    "n_mels": 100,
25	    "f_min": 0,
26	    "f_max": None,      # full bandwidth (matches Vocos mel-24khz)
27	    "power": 1.0,
28	    "clip_val": 1e-7,   # log clamp value (matches Vocos safe_log)
29	}
30
31	# ─── Segmentation config ───
32	SEG_CONFIG = {
33	    "min_duration": 3.0,   # seconds
34	    "max_duration": 12.0,  # seconds
35	    "target_duration": 8.0,  # preferred chunk size
36	    "silence_threshold_db": -40,
37	    "min_silence_len": 0.3,  # seconds
38	}
39
40
41	def detect_silence_points(waveform: np.ndarray, sr: int,
42	                          threshold_db: float = -40,
43	                          min_silence_len: float = 0.3) -> list:
44	    """Find silence points in audio for splitting."""
45	    # Convert to dB
46	    frame_length = int(0.025 * sr)  # 25ms frames
47	    hop = int(0.010 * sr)           # 10ms hop
48
49	    # RMS energy per frame
50	    n_frames = 1 + (len(waveform) - frame_length) // hop
51	    rms_db = []
52	    for i in range(n_frames):
53	        start = i * hop
54	        frame = waveform[start:start + frame_length]
55	        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
56	        rms_db.append(20 * np.log10(rms + 1e-10))
57
58	    rms_db = np.array(rms_db)
59	    is_silence = rms_db < threshold_db
60
61	    # Find contiguous silence regions
62	    silence_points = []
63	    min_frames = int(min_silence_len / 0.010)
64
65	    in_silence = False
66	    silence_start = 0
67	    for i, s in enumerate(is_silence):
68	        if s and not in_silence:
69	            silence_start = i
70	            in_silence = True
71	        elif not s and in_silence:
72	            if i - silence_start >= min_frames:
73	                # Midpoint of silence region in samples
74	                mid_frame = (silence_start + i) // 2
75	                mid_sample = mid_frame * hop
76	                silence_points.append(mid_sample)
77	            in_silence = False
78
79	    return silence_points
80
81
82	def segment_audio(waveform: np.ndarray, sr: int, text: str, config: dict) -> list:
83	    """Segment a long audio clip into shorter chunks at silence boundaries."""
84	    total_duration = len(waveform) / sr
85
86	    # If already short enough, return as-is
87	    if total_duration <= config["max_duration"]:
88	        if total_duration >= config["min_duration"]:
89	            return [(waveform, text)]
90	        else:
91	            return []  # Too short
92
93	    # Find silence points for natural splitting
94	    silence_points = detect_silence_points(
95	        waveform, sr,
96	        threshold_db=config["silence_threshold_db"],
97	        min_silence_len=config["min_silence_len"]
98	    )
99
100	    # Add start and end


Now let me check the inference_engine.py to see the production inference pipeline:


1	"""
2	Production inference engine for LeWM TTS.
3	Features:
4	  - Batched inference (many concurrent requests in parallel)
5	  - KV-cache (O(1) per step instead of O(n))
6	  - FP16
7	  - Audio prompt support (for multi-speaker / voice cloning)
8	  - Continuous batching (add/remove requests dynamically)
9	  - Griffin-Lim + Vocos vocoder pipeline
10	"""
11
12	import torch
13	import torch.nn.functional as F
14	import torchaudio
15	import numpy as np
16	import time
17	from dataclasses import dataclass, field
18	from typing import Optional
19	from pathlib import Path
20
21	from model import LeWMTTS
22
23
24	# ─── Vocoder ────────────────────────────────────────────────────────────────
25
26	class Vocoder:
27	    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""
28
29	    def __init__(self):
30	        from vocos import Vocos
31	        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
32	        self.vocos.eval()
33
34	    def __call__(self, mel):
35	        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
36	        with torch.no_grad():
37	            # Vocos decode handles batched input
38	            waveforms = []
39	            for i in range(mel.shape[0]):
40	                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
41	                waveforms.append(wav.squeeze(0).numpy())
42	        return waveforms
43
44
45	# ─── Request / Result ───────────────────────────────────────────────────────
46
47	@dataclass
48	class TTSRequest:
49	    text: str
50	    max_steps: int = 300
51	    temperature: float = 0.0
52	    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning
53
54
55	@dataclass
56	class TTSResult:
57	    waveform: np.ndarray
58	    sample_rate: int = 24000
59	    steps_taken: int = 0
60	    generation_time: float = 0.0
61
62	    def save(self, path):
63	        wav = torch.from_numpy(self.waveform).unsqueeze(0)
64	        torchaudio.save(path, wav, self.sample_rate)
65
66	    @property
67	    def duration(self):
68	        return len(self.waveform) / self.sample_rate
69
70
71	# ─── Inference Engine ───────────────────────────────────────────────────────
72
73	class InferenceEngine:
74	    """
75	    High-throughput batched TTS inference engine.
76
77	    Usage:
78	        engine = InferenceEngine("checkpoint.pt")
79
80	        # Single request
81	        result = engine.synthesize("नमस्ते भारत")
82
83	        # Batched (high throughput)
84	        results = engine.synthesize_batch(["text1", "text2", ...])
85
86	        # With voice prompt
87	        result = engine.synthesize("text", prompt_audio="ref.wav")
88
89	        # Benchmark
90	        engine.benchmark()
91	    """
92
93	    def __init__(self, checkpoint_path, device="cuda", dtype=torch.float16):
94	        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
95	        self.dtype = dtype
96
97	        # Load model
98	        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
99	        self.config = ckpt["config"]
100	        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
101	        self.model.load_state_dict(ckpt["model"])
102	        self.model.eval()
103
104	        self.d_model = self.config["d_model"]
105
106	        # Vocoder (runs on CPU)
107	        self.vocoder = Vocoder()
108
109	        print(f"Engine ready: {checkpoint_path}")
110	        print(f"  Device: {self.device}, Dtype: {dtype}")
111	        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")
112
113	    # ─── Text encoding ──────────────────────────────────────────────────
114
115	    def _encode_texts(self, texts):
116	        """Encode list of texts → padded embeddings + mask.
117	        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
118	        """
119	        # Byte-level tokenization
120	        token_lists = [list(t.encode("utf-8")) for t in texts]
121	        max_len = max(len(t) for t in token_lists)
122
123	        # Pad to same length
124	        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
125	        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)
126
127	        for i, tokens in enumerate(token_lists):
128	            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
129	            mask[i, :len(tokens)] = False
130
131	        with torch.no_grad():
132	            text_emb = self.model.text_encoder(padded, text_mask=mask)
133
134	        return text_emb, mask
135
136	    # ─── Audio prompt encoding ──────────────────────────────────────────
137
138	    def _encode_prompt(self, prompt_mel):
139	        """Encode prompt mel → embeddings for AR seeding.
140	        prompt_mel: [1, 80, T]
141	        Returns: prompt_embs [1, T_down, d]
142	        """
143	        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
144	        with torch.no_grad():
145	            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
146	        return embs
147
148	    def load_prompt_audio(self, audio_path, max_seconds=3.0):
149	        """Load audio file → mel for prompting.
150	        Returns: mel [1, 80, T]
151	        """
152	        wav, sr = torchaudio.load(audio_path)
153	        if sr != 24000:
154	            wav = torchaudio.functional.resample(wav, sr, 24000)
155	        # Trim to max_seconds
156	        max_samples = int(max_seconds * 24000)
157	        wav = wav[:1, :max_samples]  # mono, trimmed
158
159	        mel_transform = torchaudio.transforms.MelSpectrogram(
160	            sample_rate=24000, n_fft=1024, hop_length=256,
161	            n_mels=100, power=1.0,
162	        )
163	        mel = mel_transform(wav)
164	        mel = torch.log(mel.clamp(min=1e-7))
165	        return mel  # [1, 80, T]
166
167	    # ─── Core AR generation (batched) ───────────────────────────────────
168
169	    def _generate_batch(self, text_emb, text_mask, batch_size,
170	                        max_steps=300, temperature=0.0,
171	                        prompt_embs=None):
172	        """
173	        Core batched AR generation with KV-cache.
174
175	        Args:
176	            text_emb: [B, T_text, d]
177	            text_mask: [B, T_text]
178	            batch_size: int
179	            max_steps: int
180	            temperature: float
181	            prompt_embs: list of [1, T_prompt, d] or None per request
182
183	        Returns:
184	            all_embs: [B, T_total, d]
185	            steps_taken: list of int per request
186	        """
187	        d = self.d_model
188	        cache = self.model.init_ar_cache()
189
190	        # Track per-request state
191	        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
192	        steps_taken = [0] * batch_size
193	        all_emb_lists = [[] for _ in range(batch_size)]
194	        prev_norms = [[] for _ in range(batch_size)]
195
196	        # Handle prompts: seed cache with prompt embeddings
197	        if prompt_embs is not None:
198	            max_prompt_len = max(
199	                (e.shape[1] for e in prompt_embs if e is not None), default=0
200	            )


Perfect! Now let me check the v2 directory for any additional insights:


total 60
drwxrwxr-x  5 ubuntu ubuntu 4096 Mar 23 22:05 .
drwxrwxr-x 29 ubuntu ubuntu 4096 Mar 24 19:41 ..
-rw-rw-r--  1 ubuntu ubuntu 2190 Mar 23 21:35 dataset.py
-rw-rw-r--  1 ubuntu ubuntu 4744 Mar 23 22:02 inference.py
-rw-rw-r--  1 ubuntu ubuntu 8702 Mar 23 22:04 model.py
drwxrwxr-x  3 ubuntu ubuntu 4096 Mar 25 02:08 output
-rw-rw-r--  1 ubuntu ubuntu 6931 Mar 23 21:40 preprocess.py
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 25 02:08 processed_data
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 23 22:05 __pycache__
-rw-rw-r--  1 ubuntu ubuntu 7228 Mar 23 22:05 train.py


Let me check the v2 model to see if there are differences:


1	"""
2	LeWM TTS v2 — DAC-based JEPA Text-to-Speech
3
4	Architecture:
5	  - DAC encoder (frozen): audio → 1024-dim continuous latents @ 75Hz
6	  - Linear proj: 1024 → 256 (d_model)
7	  - TextEncoder: byte-level transformer
8	  - JEPAPredictor: causal transformer decoder with cross-attention to text
9	  - Linear proj: 256 → 1024 (back to DAC space)
10	  - DAC decoder (frozen): latents → waveform
11
12	  Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL)
13	  NO mel decoder needed — DAC handles audio reconstruction.
14	"""
15
16	import math
17	import torch
18	import torch.nn as nn
19	import torch.nn.functional as F
20
21
22	class SinusoidalPositionalEncoding(nn.Module):
23	    def __init__(self, d_model, max_len=8192):
24	        super().__init__()
25	        pe = torch.zeros(max_len, d_model)
26	        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
27	        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
28	        pe[:, 0::2] = torch.sin(position * div_term)
29	        pe[:, 1::2] = torch.cos(position * div_term)
30	        self.register_buffer("pe", pe.unsqueeze(0))
31
32	    def forward(self, x):
33	        return x + self.pe[:, :x.shape[1]]
34
35
36	class TextEncoder(nn.Module):
37	    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
38	        super().__init__()
39	        self.d_model = d_model
40	        self.char_embed = nn.Embedding(vocab_size, d_model)
41	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
42	        self.dropout = nn.Dropout(dropout)
43
44	        encoder_layer = nn.TransformerEncoderLayer(
45	            d_model=d_model, nhead=nhead,
46	            dim_feedforward=d_model * 4,
47	            dropout=dropout, batch_first=True, activation="gelu",
48	        )
49	        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
50	        self.proj = nn.Linear(d_model, d_model)
51
52	    def forward(self, text_tokens, text_mask=None):
53	        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
54	        x = self.pos_embed(x)
55	        x = self.dropout(x)
56	        x = self.transformer(x, src_key_padding_mask=text_mask)
57	        return self.proj(x)
58
59
60	class JEPAPredictor(nn.Module):
61	    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
62	        super().__init__()
63	        self.audio_input_proj = nn.Linear(d_model, d_model)
64	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
65	        self.dropout = nn.Dropout(dropout)
66
67	        decoder_layer = nn.TransformerDecoderLayer(
68	            d_model=d_model, nhead=nhead,
69	            dim_feedforward=d_model * 4,
70	            dropout=dropout, batch_first=True, activation="gelu",
71	        )
72	        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
73	        self.output_proj = nn.Sequential(
74	            nn.Linear(d_model, d_model),
75	            nn.GELU(),
76	            nn.Linear(d_model, d_model),
77	        )
78
79	    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
80	        x = self.audio_input_proj(audio_emb)
81	        x = self.pos_embed(x)
82	        x = self.dropout(x)
83
84	        T = x.shape[1]
85	        causal_mask = torch.triu(
86	            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
87	        )
88
89	        predicted = self.transformer(
90	            x, text_emb,
91	            tgt_mask=causal_mask,
92	            tgt_key_padding_mask=audio_mask,
93	            memory_key_padding_mask=text_mask,
94	        )
95	        return self.output_proj(predicted)
96
97
98	class LeWMTTS(nn.Module):
99	    """
100	    LeWM TTS v2 with DAC latents.
101
102	    Training:
103	      - DAC encodes audio → 1024-dim latents (frozen)
104	      - Project to d_model, predict next embedding, KL regularize
105	    Inference:
106	      - Text → predict DAC latents autoregressively → DAC decode → waveform
107	    """
108
109	    def __init__(self, config):
110	        super().__init__()
111	        d_model = config.get("d_model", 256)
112	        nhead = config.get("nhead", 4)
113	        dac_dim = config.get("dac_dim", 1024)
114	        text_vocab = config.get("text_vocab_size", 256)
115	        text_layers = config.get("text_encoder_layers", 4)
116	        predictor_layers = config.get("predictor_layers", 6)
117	        dropout = config.get("dropout", 0.1)
118	        self.kl_weight = config.get("kl_weight", 0.1)
119	        self.recon_weight = config.get("recon_weight", 10.0)
120
121	        # DAC space ↔ model space projections
122	        self.dac_in_proj = nn.Linear(dac_dim, d_model)
123	        self.dac_out_proj = nn.Linear(d_model, dac_dim)
124
125	        # Gaussian regularizer projections
126	        self.proj_mu = nn.Linear(d_model, d_model)
127	        self.proj_logvar = nn.Linear(d_model, d_model)
128
129	        self.text_encoder = TextEncoder(
130	            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
131	            num_layers=text_layers, dropout=dropout,
132	        )
133	        self.predictor = JEPAPredictor(
134	            d_model=d_model, nhead=nhead,
135	            num_layers=predictor_layers, dropout=dropout,
136	        )
137
138	        self.config = config
139
140	    def forward(self, dac_latents, text_tokens, latent_mask=None, text_mask=None):
141	        """
142	        Args:
143	            dac_latents: [B, T, 1024] — continuous DAC encoder output
144	            text_tokens: [B, T_text] — byte-level text tokens
145	            latent_mask: [B, T] bool, True = padding
146	            text_mask: [B, T_text] bool, True = padding
147	        """
148	        # Project DAC latents to model space
149	        h = self.dac_in_proj(dac_latents)  # [B, T, d_model]
150
151	        # Gaussian regularizer: project to mu/logvar
152	        mu = self.proj_mu(h)
153	        logvar = self.proj_logvar(h)
154
155	        # Reparameterize
156	        if self.training:
157	            std = torch.exp(0.5 * logvar)
158	            z = mu + torch.randn_like(std) * std
159	        else:
160	            z = mu
161
162	        # Encode text
163	        text_emb = self.text_encoder(text_tokens, text_mask)
164
165	        # Predict next embedding
166	        input_emb = z[:, :-1]
167	        target_emb = z[:, 1:]
168
169	        if latent_mask is not None:
170	            pred_mask = latent_mask[:, :-1]
171	        else:
172	            pred_mask = None
173
174	        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
175
176	        # Loss 1: Next-embedding prediction (MSE)
177	        if pred_mask is not None:
178	            valid = (~pred_mask).unsqueeze(-1)
179	            prediction_loss = F.mse_loss(
180	                predicted * valid, target_emb * valid, reduction="sum"
181	            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
182	        else:
183	            prediction_loss = F.mse_loss(predicted, target_emb)
184
185	        # Loss 2: KL(N(mu, sigma²) || N(0, I))
186	        if latent_mask is not None:
187	            valid = (~latent_mask).unsqueeze(-1).float()
188	            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
189	            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
190	        else:
191	            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
192
193	        # Loss 3: Projection round-trip reconstruction
194	        # Forces dac_out_proj(dac_in_proj(x)) ≈ x so DAC can decode our outputs
195	        dac_recon = self.dac_out_proj(mu)  # Use mu (not noisy z) for cleaner target
196	        if latent_mask is not None:
197	            valid = (~latent_mask).unsqueeze(-1).float()
198	            recon_loss = F.mse_loss(
199	                dac_recon * valid, dac_latents * valid, reduction="sum"
200	            ) / (valid.sum() * dac_latents.shape[-1] + 1e-8)
201	        else:
202	            recon_loss = F.mse_loss(dac_recon, dac_latents)
203
204	        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss
205
206	        return {
207	            "total_loss": total_loss,
208	            "prediction_loss": prediction_loss,
209	            "kl_loss": kl_loss,
210	            "recon_loss": recon_loss,
211	        }
212
213	    def predict_next(self, audio_emb, text_emb, text_mask=None):
214	        """AR inference: predict next embedding from context."""
215	        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
216	        return predicted[:, -1:]
217
218	    def latents_to_dac(self, z):
219	        """Project model-space embeddings back to DAC latent space."""
220	        return self.dac_out_proj(z)
221
222
223	def count_parameters(model):
224	    return sum(p.numel() for p in model.parameters() if p.requires_grad)
225
226
227	def build_model(config=None):
228	    if config is None:
229	        config = {
230	            "d_model": 256,
231	            "nhead": 4,
232	            "dac_dim": 1024,
233	            "text_vocab_size": 256,
234	            "text_encoder_layers": 4,
235	            "predictor_layers": 6,
236	            "dropout": 0.1,
237	            "kl_weight": 0.1,
238	        }
239	    model = LeWMTTS(config)
240	    print(f"LeWM TTS v2 (DAC): {count_parameters(model)/1e6:.2f}M trainable parameters")
241	    return model, config
242


Excellent! Now let me create a comprehensive summary document. Let me verify a few more details about the speaker conditioning implementation:


481:        n_speakers = config.get("n_speakers", 1)
510:        # Speaker embedding (for multi-speaker)
511:        self.n_speakers = n_speakers
512:        if n_speakers > 1:
513:            self.speaker_embed = nn.Embedding(n_speakers, d_model)
515:            self.speaker_embed = None
535:    def _add_speaker(self, text_emb, speaker_id):
536:        """Add speaker embedding to text embeddings."""
537:        if self.speaker_embed is not None and speaker_id is not None:
538:            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
549:    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
559:        # Encode text + speaker conditioning
561:        text_emb = self._add_speaker(text_emb, speaker_id)


Perfect! Now let me create a comprehensive report. Let me check one more file to verify the training predictor architecture details:


grep -n "free_bits\|ema_decay\|start_emb" /home/ubuntu/lewm-tts/model.py | head -20


508:        self.ema_decay = config.get("ema_decay", 0.998)
518:        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
527:        self.free_bits = config.get("free_bits", 1.0)
547:            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
577:        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
579:        # This trains the model to predict the first audio embedding from start_emb + text
580:        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
583:        # Input noise: only on positions 1+ (real embeddings), not start_emb
594:            # Predictor input mask: start_emb (pos 0) is always unmasked
617:        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
701:            "ema_decay": 0.998,
702:            "free_bits": 1.0,


Perfect! Now I have all the information needed. Let me create a comprehensive report:

## LeWM TTS Codebase - Comprehensive Exploration Report

I've completed a thorough exploration of the LeWM TTS codebase at `/home/ubuntu/lewm-tts`. Here's a complete breakdown of all the key components:

---

### 1. MAIN MODEL ARCHITECTURE FILES

#### **Primary Model: `/home/ubuntu/lewm-tts/model.py`** (708 lines)
The core JEPA-based TTS implementation with 4 main components:

**TextEncoder:**
- Byte-level character transformer for Hindi/UTF-8 text
- Config: 256d embedding, 4 attention heads, 4 transformer layers
- Character embedding + sinusoidal positional encoding
- Linear projection output

**AudioEncoder:**
- 1D CNN downsampler (stride-2 or stride-4) + Transformer contextualization
- Input: mel spectrogram [B, n_mels, T_mel]
- 1D convolutions with GELU activation for temporal compression
- Sinusoidal positional encoding + 4 transformer encoder layers
- **Gaussian latent space** with reparameterization:
  - `proj_mu`: projects to mean (μ)
  - `proj_logvar`: projects to log-variance (σ)
  - During training: z = μ + ε·σ (reparameterization trick)
  - During inference: z = μ (deterministic)

**JEPAPredictor:**
- Causal transformer decoder with **cross-attention to text**
- 6 transformer layers with:
  - Self-attention over audio embeddings (causal masking)
  - Cross-attention to text embeddings
  - Feed-forward networks
- Output projection to predict next audio embedding
- **KV-cache support** for efficient autoregressive inference (lines 209-323):
  - `_cached_mha()`: multi-head attention with cached K,V
  - `_cached_cross_attn()`: cross-attention with pre-computed text K,V
  - `init_cache()`: initialize empty KV cache
  - `inference_step()`: process one token in O(1) time using cached attention

**MelDecoder:**
- Reconstructs mel spectrograms from latent embeddings
- Architecture:
  - Linear projection to 512d (2× d_model)
  - Transposed convolution upsample (2×2 = 4×)
  - ResConvBlock residual blocks for sharpness
  - Final conv refinement → n_mels output
- Output: [B, n_mels, T_up]

**MultiResolutionSpectralLoss:**
- STFT loss at 3 resolutions: (64/16/64), (128/32/128), (256/64/256)
- For each resolution: spectral convergence + log magnitude loss
- Forces sharp, well-defined mel spectral structure

**Full LeWMTTS Model:**
- **EMA target encoder** (core JEPA component):
  - `ema_audio_encoder`: exponential moving average copy of online encoder
  - `ema_decay`: 0.998 (updated each training step)
  - Provides slow-moving, stable prediction targets
  - Prevents representation collapse
- **Learnable start embedding**: [1, 1, d_model] to seed AR generation
- **Speaker conditioning** (optional multi-speaker):
  - `speaker_embed`: embedding table (if n_speakers > 1)
  - Added to text embeddings via `_add_speaker()`
- **4 Loss functions** (lines 605-654):
  1. **Prediction Loss (MSE)**: MSE(predicted[t], ema_target[t+1])
  2. **KL Loss**: Free-bits KL divergence with annealing
     - `kl_free = max(0, kl_per_dim - free_bits)` (prevents collapse)
  3. **Reconstruction Loss (L1)**: mel_decoder output vs ground truth mel
  4. **Spectral Loss (STFT)**: multi-resolution spectral convergence

---

#### **V5 Model: `/home/ubuntu/lewm-tts/model_v5.py`** (463 lines)
Evolved codec-based variant using EnCodec instead of mel:

**Key Differences:**
- Encodes audio with **EnCodec (128-dim embeddings @ 75Hz)**
- Simple linear projections instead of convolutional encoder/decoder
- **Duration predictor** for text-to-audio alignment (FastSpeech-style)
- Frame-level text conditioning (text expanded then added to audio input)

**DurationPredictor:**
- Conv1d layers with LayerNorm and dropout
- Predicts log-duration per text character
- `compute_uniform_durations()`: spreads audio frames evenly across text

**JEPAPredictor (v5):**
- Uses **TransformerEncoder** (self-attention only)
- Text embeddings expanded to frame-level and **summed with audio input**
- Cannot ignore text—it's baked into the input
- Backward-compatible cross-attention interface

**LeWMTTSv5:**
- `proj_in`: codec_dim (128) → d_model (256)
- `proj_out`: d_model (256) → codec_dim (128)
- EMA on `proj_in` (provides targets)
- **3 Loss functions**:
  1. Prediction loss (MSE)
  2. Roundtrip loss (proj_in → proj_out → match original codec embedding)
  3. Duration loss (MSE on log-durations)

---

#### **V2 Model: `/home/ubuntu/lewm-tts/v2/model.py`** (242 lines)
Earlier DAC-based variant:

- DAC encoder (1024-dim latents @ 75Hz, frozen)
- Simple projection: 1024 ↔ 256
- Gaussian regularizer projections (mu/logvar)
- Cross-attention predictor
- No mel decoder—DAC handles reconstruction
- Losses: prediction (MSE) + KL + roundtrip reconstruction

---

### 2. SPEAKER CONDITIONING IMPLEMENTATION

**Location:** `/home/ubuntu/lewm-tts/model.py`, lines 510-540

**Method: Additive Embedding**
```python
self.speaker_embed = nn.Embedding(n_speakers, d_model)  # if n_speakers > 1

def _add_speaker(self, text_emb, speaker_id):
    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
```

- Speaker ID is converted to d_model-dimensional embedding
- Added to text embeddings (broadcast across time dimension)
- Called in forward pass before passing to predictor
- Currently designed but **not yet trained** (mentioned in README as "planned")

---

### 3. TRAINING LOOP & LOSS FUNCTIONS

**Main Training: `/home/ubuntu/lewm-tts/train.py`** (246 lines)

**Loss Weights (configurable):**
```python
pred_weight: 10.0        # Next-embedding prediction
kl_weight: 0.05-0.1      # Gaussian regularizer
recon_weight: 1.0        # Mel reconstruction
spectral_weight: 0.5     # Multi-resolution spectral
```

**Total Loss:**
```python
loss = (pred_weight × prediction_loss
        + kl_weight × kl_loss
        + recon_weight × recon_loss
        + spectral_weight × spectral_loss)
```

**Annealing Strategies:**
- **KL Annealing**: ramps from 0 to 1 over first 20% of training
- **Input Noise Annealing**: ramps from 0 to 0.2 over first 30% of training
  - Bridges train/inference gap by simulating imperfect AR predictions
  - Only applied to positions 1+ (not start_emb)

**Free-Bits KL:**
```python
kl_free = torch.clamp(kl_per_dim - free_bits, min=0.0)  # free_bits=1.0
```
- Allows KL up to threshold without penalty
- Prevents both collapse (too low KL) and explosion (too high KL)

**Optimizer & Scheduler:**
- AdamW: lr=1e-4, betas=(0.9, 0.98), weight_decay=0.01
- Warmup: min(4000 steps, 10% of total)
- Cosine annealing decay over remaining steps
- Gradient clipping: 1.0

**EMA Update:**
```python
model.update_ema()  # After each optimizer step
```

**V5 Training: `/home/ubuntu/lewm-tts/train_v5.py`** (similar structure)
- Codec embedding input instead of mel
- 3 losses: prediction, roundtrip, duration
- Input noise annealing only

---

### 4. JEPA ENCODER/PREDICTOR/TARGET ENCODER INTERACTION

**Core JEPA Loop (lines 542-612 in model.py):**

1. **Online Encoder** processes mel spectrogram:
   ```python
   z, mu, logvar = self.audio_encoder(mel, mel_mask)  # [B, T_down, d]
   ```

2. **EMA Target Encoder** (no gradients, always eval):
   ```python
   with torch.no_grad():
       self.ema_audio_encoder.eval()
       _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
       self.ema_audio_encoder.train()
       target_emb = mu_ema  # [B, T_down, d]
   ```

3. **Prepare Input Sequence** (shifted right):
   ```python
   start = self.start_emb.expand(B, -1, -1)        # [B, 1, d]
   input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
   # Input:  [start_emb, z[0], z[1], ..., z[T-2]]
   # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]]
   ```

4. **Optionally Add Input Noise** (training only):
   ```python
   if self.training and self.input_noise > 0:
       noise[:, 1:] = randn(...) * self.input_noise
       input_emb = input_emb + noise
   ```

5. **Predictor** predicts next embeddings:
   ```python
   predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   # predicted[t] should match target_emb[t]
   ```

6. **EMA Update**:
   ```python
   for p_online, p_ema in zip(audio_encoder.parameters(), ema_encoder.parameters()):
       p_ema = p_ema * 0.998 + p_online * 0.002
   ```

**Why This Design:**
- **Prevents collapse**: EMA encoder provides slowly-moving targets
- **Prediction learning**: Predictor learns to extrapolate from text context
- **Stability**: No gradient feedback to targets; prevents training instability

---

### 5. CONFIG FILES & MODEL DIMENSIONS

**Main Config (output_v3/config.json):**
```json
{
  "d_model": 256,
  "nhead": 4,
  "n_mels": 100,
  "text_vocab_size": 256,
  "text_encoder_layers": 4,
  "audio_encoder_layers": 4,
  "predictor_layers": 6,
  "downsample_factor": 4,
  "dropout": 0.1,
  "kl_weight": 0.1,
  "recon_weight": 1.0,
  "spectral_weight": 0.5,
  "pred_weight": 10.0,
  "ema_decay": 0.998,
  "free_bits": 1.0
}
```

**V5 Config (output_v5/config.json):**
```json
{
  "d_model": 256,
  "nhead": 4,
  "codec_dim": 128,
  "text_vocab_size": 256,
  "text_encoder_layers": 4,
  "predictor_layers": 6,
  "dropout": 0.1,
  "pred_weight": 10.0,
  "roundtrip_weight": 1.0,
  "ema_decay": 0.998,
  "input_noise": 0.0
}
```

**Key Dimensions:**
- d_model: 256 (embedding dimension)
- nhead: 4 (attention heads, head_dim=64)
- text_vocab_size: 256 (byte-level UTF-8 encoding, 0-255)
- n_mels: 100 (mel bins, matches Vocos)
- downsample_factor: 4 (CNN stride in audio encoder)
- Total parameters: ~24M (v3), similar for v5

---

### 6. INFERENCE & SYNTHESIS PIPELINE

**Basic Inference: `/home/ubuntu/lewm-tts/inference.py`** (151 lines)

**Class: LeWMTTSInference**

**Methods:**

1. **`synthesize(text, max_steps=300, temperature=0.0)`** (lines 48-82)
   - Text → byte-level tokens
   - Encode text: TextEncoder → [B, T_text, d]
   - Use learned start_emb to seed
   - **Autoregressive loop with KV-cache**:
     ```python
     all_embeddings = [start_emb]
     cache = model.init_ar_cache()

     for step in range(max_steps):
         next_emb, cache = model.predict_next_cached(
             next_emb, text_emb, step, cache
         )
         all_embeddings.append(next_emb)

         # Stop early if embeddings stabilize
         if embedding_norm_variance < threshold:
             break
     ```
   - Concatenate embeddings: [B, T_total, d]
   - MelDecoder → mel [B, 100, T]
   - Vocos vocoder → waveform

2. **`reconstruct(mel_path)`** (lines 85-93)
   - Encode real mel
   - MelDecoder reconstruction
   - Vocos decode

3. **`reconstruct_with_prediction(mel_path)`** (lines 96-110)
   - Encode real mel → z
   - Predictor with dummy text (teacher-forced)
   - Decode predicted embeddings

---

**V5 Inference: `/home/ubuntu/lewm-tts/inference_v5.py`** (263 lines)

**Methods:**

1. **`synthesize(text, duration_scale=1.0, temperature=0.0)`** (lines 40-87)
   - Text → TextEncoder → [B, T_text, d]
   - **DurationPredictor** → log_durations → durations
   - `length_regulate()`: expand text embeddings by durations
   - Autoregressive loop (similar to v1)
   - proj_out → codec embeddings
   - EnCodec decoder → waveform

2. **`synthesize_prompted(text, prompt_audio, ...)`** (lines 90-160)
   - Encode prompt audio with EnCodec
   - Process prompt frames through predictor cache
   - Continue AR generation from prompt
   - Allows voice cloning by providing reference audio

3. **`reconstruct(audio_path)`** (lines 163-173)
   - Tests roundtrip: encode → proj_in → proj_out → decode

4. **`reconstruct_with_prediction(audio_path, text)`** (lines 187-220)
   - Encode audio → z
   - Expand text uniformly to audio length
   - Predictor with teacher forcing
   - Decode

---

**Production Inference Engine: `/home/ubuntu/lewm-tts/inference_engine.py`** (partial view)

**Features:**
- Batched inference (concurrent requests)
- FP16 support for GPU efficiency
- KV-cache for O(1) per-step inference
- Audio prompt encoding for voice cloning
- Vocos vocoder integration
- Benchmarking utilities

**Performance (A100 80GB):**
- Single request: 9× real-time
- Batch size 64: 359× real-time throughput

---

### 7. DATASET & PREPROCESSING

**Dataset: `/home/ubuntu/lewm-tts/dataset.py`** (98 lines)

**TTSDataset:**
- Loads from manifest.json (list of entries with mel_path, text)
- Max length filtering: 1200 mel frames, 512 text tokens
- Byte-level text tokenization (UTF-8 encoding)
- Returns: {mel [n_mels, T], text_tokens [T_text]}

**Collate Function:**
- Pads mel specs to max length along time axis
- Pads text tokens to max length
- Creates boolean masks (True = padding)
- Returns: mel [B, n_mels, T_max], text_tokens [B, T_text_max], mel_mask, text_mask

**V5 Dataset: `/home/ubuntu/lewm-tts/dataset_v5.py`** (68 lines)

**CodecDataset:**
- Loads precomputed EnCodec embeddings [128, T]
- Byte-level text tokens
- Max 900 codec frames
- Collate: pad codec embeddings [B, 128, T_max], text tokens [B, T_text_max]

---

**Preprocessing: `/home/ubuntu/lewm-tts/preprocess.py`** (100+ lines)

**Mel Config:**
```python
sample_rate: 24000
n_fft: 1024
hop_length: 256
n_mels: 100
f_min: 0
f_max: None  # full bandwidth
power: 1.0
clip_val: 1e-7  # log clamp (matches Vocos)
```

**Segmentation Config:**
- Min duration: 3 seconds
- Max duration: 12 seconds
- Target chunk: 8 seconds
- Silence threshold: -40 dB
- Min silence: 0.3 seconds

**Process:**
1. Load audio from HuggingFace dataset
2. Detect silence points (RMS-based)
3. Segment into 3-12s chunks at silence boundaries
4. Extract mel spectrograms
5. Save preprocessed dataset + manifest.json

**Dataset Stats:**
- Source: 20 hours single-speaker Hindi audio
- Processed: 7,009 segments, 17.88 hours
- Mel extraction: matches Vocos `mel-24khz` exactly

---

### 8. DIRECTORY STRUCTURE

```
/home/ubuntu/lewm-tts/
├── model.py                    # Main v1 model (mel-based)
├── model_v5.py                 # V5 model (codec-based)
├── train.py                    # Training loop v1
├── train_v5.py                 # Training loop v5
├── inference.py                # Single-request inference v1
├── inference_v5.py             # Single-request inference v5
├── inference_engine.py         # Production batched inference
├── dataset.py                  # PyTorch dataset (mel-based)
├── dataset_v5.py               # PyTorch dataset (codec-based)
├── preprocess.py               # Data preprocessing
├── README.md                    # Documentation

├── v2/                         # Earlier DAC-based variant
│   ├── model.py
│   ├── train.py
│   ├── inference.py
│   ├── dataset.py
│   ├── preprocess.py
│   └── output/
│       └── checkpoints/

├── processed_data/             # Mel-based dataset
│   ├── manifest.json
│   └── config.json

├── processed_data_codec/       # Codec-based dataset
│   └── manifest.json

├── output_v3/                  # Latest mel-based checkpoint
│   ├── checkpoints/
│   ├── logs/
│   └── config.json

├── output_v5/                  # Latest codec-based checkpoint
│   ├── checkpoints/
│   ├── logs/
│   └── config.json

├── samples_v5/                 # Generated audio samples
├── samples_v5_pred/
├── samples_v5_prompted/        # Prompted synthesis outputs
└── [other output dirs...]
```

---

### 9. ARCHITECTURE SUMMARY DIAGRAM

```
TEXT INPUT (UTF-8 bytes)
         ↓
    TextEncoder [4 layers, 256d, 4 heads]
         ↓
   text_emb [B, T_text, 256d]
         ↓ (+ optional speaker embedding)

                            MEL INPUT [B, 100, T_mel]
                                      ↓
                         AudioEncoder (1D CNN + Transformer)
                         1. Conv1d stride-4 downsample
                         2. Transformer [4 layers, 256d, 4 heads]
                         3. Gaussian projections (mu/logvar)
                                      ↓
                         z [B, T_down, 256d] (sampled)
                         mu [B, T_down, 256d] (mean)
                         logvar [B, T_down, 256d] (log-variance)

                         EMA AudioEncoder (same, no gradient, 0.998 decay)
                                      ↓
                         target_emb [B, T_down, 256d]

    [start_emb] + z[:-1] ──→ JEPAPredictor ←── text_emb
                              6 layers:
                              - Causal self-attn
                              - Cross-attn to text
                              - FFN
                                      ↓
                         predicted [B, T_down, 256d]
                                      ↓
                         MelDecoder (linear + transpose conv + ResBlocks)
                                      ↓
                         mel_recon [B, 100, T_mel]
                                      ↓
                         Vocos Vocoder (24khz)
                                      ↓
                         waveform [24000 hz]
```

---

### 10. KEY TRAINING INSIGHTS

1. **Representation Learning**: The EMA target encoder provides stable targets that prevent the online encoder from collapsing to a meaningless representation.

2. **Text Conditioning**: Cross-attention from predictor to text embeddings allows the model to learn text-aligned predictions, grounding the latent space in linguistic structure.

3. **Multi-Scale Losses**: Combination of MSE (prediction), KL (regularization), L1 (reconstruction), and STFT (spectral detail) ensures well-formed latent representations and sharp mel output.

4. **Free-Bits Regularization**: Prevents both KL collapse (model ignores text) and excessive KL divergence (latent space becomes too stochastic).

5. **Input Noise & Annealing**: Simulates the AR inference distribution during training, bridging the train/test gap where errors accumulate.

6. **Causal Masking**: The predictor uses causal attention, ensuring it only attends to past audio and current text context—enforcing autoregressiveness.

7. **Vocab Design**: Byte-level encoding (256 tokens) handles any Unicode text (Hindi/English/etc.) without needing language-specific tokenizers.

---

### 11. INFERENCE OPTIMIZATION

- **KV-Cache**: Reduces per-step complexity from O(T) to O(1) by caching previous attention keys/values
- **FP16**: Reduces memory footprint and increases throughput
- **Batch Processing**: Concurrent inference on multiple requests with shared compute
- **Early Stopping**: Monitors embedding norm variance to stop AR generation when stable

All files are located in `/home/ubuntu/lewm-tts/` with subdirectories for different training variants (v2, v3, v5) and output checkpoints.