Thoroughly explore the LeWM TTS codebase in /home/ubuntu. Find:

1. The main model architecture files — encoder, decoder, predictor, any JEPA-related components
2. How speaker conditioning is implemented (speaker embeddings, FiLM, cross-attention, etc.)
3. The training loop and loss functions
4. How the JEPA encoder/predictor/target encoder interact
5. Any config files showing model dimensions, layers, etc.
6. The inference/synthesis pipeline

Look for directories like lewm*, tts*, model*, src*, etc. Check for Python files with model definitions, training scripts, config files (yaml, json, toml). Be very thorough — check all subdirectories.


I'll conduct a thorough exploration of the LeWM TTS codebase. Let me start by identifying the directory structure and then systematically search for the relevant components.


Bash

find /home/ubuntu -maxdepth 3 -type d -name "*lewm*" -o -type d -name "*tts*" -o -type d -name "*TTS*" 2>/dev/null | head -20


/home/ubuntu/veenaModal/external/sparktts
/home/ubuntu/tts
/home/ubuntu/tts/tts
/home/ubuntu/kanitts-2-dataset-pipeline
/home/ubuntu/.cursor/projects/home-ubuntu-tts
/home/ubuntu/.cache/claude-cli-nodejs/-home-ubuntu-kani-tts-inference
/home/ubuntu/lewm-tts
/home/ubuntu/xtts_samples
/home/ubuntu/Irodori-TTS
/home/ubuntu/Irodori-TTS/irodori_tts
/home/ubuntu/training/maya-tts-llama32-1b-pretraining
/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts
/home/ubuntu/training/checkpoints/llama3.2-1b-tts
/home/ubuntu/indic_tts_outputs
/home/ubuntu/LuxTTS


Great! I found `/home/ubuntu/lewm-tts`. Let me explore its structure comprehensively.


ls -la /home/ubuntu/lewm-tts/ && echo "=== Subdirectories ===" && find /home/ubuntu/lewm-tts -type d | head -20


total 17320
drwxrwxr-x 29 ubuntu ubuntu   4096 Mar 24 19:41 .
drwxr-x--x 82 ubuntu ubuntu  12288 Mar 30 07:43 ..
-rw-rw-r--  1 ubuntu ubuntu   2891 Mar 23 20:51 dataset.py
-rw-rw-r--  1 ubuntu ubuntu   2360 Mar 24 15:19 dataset_v5.py
-rw-rw-r--  1 ubuntu ubuntu  19643 Mar 24 14:09 inference_engine.py
-rw-rw-r--  1 ubuntu ubuntu   5586 Mar 24 14:02 inference.py
-rw-rw-r--  1 ubuntu ubuntu  11279 Mar 24 19:25 inference_v5.py
-rw-rw-r--  1 ubuntu ubuntu  28786 Mar 24 14:09 model.py
-rw-rw-r--  1 ubuntu ubuntu  19222 Mar 24 18:48 model_v5.py
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:06 output
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:06 output_v3
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:06 output_v3b
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:06 output_v3c
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:06 output_v4
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:08 output_v5
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:08 output_v5_pred
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:08 output_v5_pred_dur
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:08 output_v5_pred_mask
-rw-rw-r--  1 ubuntu ubuntu   2785 Mar 24 15:18 preprocess_codec.py
-rw-rw-r--  1 ubuntu ubuntu   9475 Mar 24 04:40 preprocess.py
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 25 02:08 processed_data
drwxrwxr-x  3 ubuntu ubuntu   4096 Mar 25 02:08 processed_data_100mel
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 25 02:08 processed_data_codec
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 19:25 __pycache__
-rw-rw-r--  1 ubuntu ubuntu   5356 Mar 24 06:01 README.md
-rwxrwxr-x  1 ubuntu ubuntu    697 Mar 23 21:57 run_v1.sh
-rw-rw-r--  1 ubuntu ubuntu 613966 Mar 24 07:01 sample_ar_prompted_v3b.wav
-rw-rw-r--  1 ubuntu ubuntu 613966 Mar 24 06:50 sample_ar_prompted.wav
-rw-rw-r--  1 ubuntu ubuntu 409166 Mar 24 08:12 sample_ar_v3c.wav
-rw-rw-r--  1 ubuntu ubuntu 456782 Mar 24 06:46 sample_gt_0.wav
-rw-rw-r--  1 ubuntu ubuntu 359502 Mar 24 06:46 sample_gt_100.wav
-rw-rw-r--  1 ubuntu ubuntu 384078 Mar 24 06:46 sample_gt_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456782 Mar 24 06:34 sample_gt_v3b.wav
-rw-rw-r--  1 ubuntu ubuntu 456782 Mar 24 14:38 sample_gt_v4_0.wav
-rw-rw-r--  1 ubuntu ubuntu 359502 Mar 24 14:38 sample_gt_v4_100.wav
-rw-rw-r--  1 ubuntu ubuntu 384078 Mar 24 14:38 sample_gt_v4_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 06:17 sample_recon_ep43_0.wav
-rw-rw-r--  1 ubuntu ubuntu 357966 Mar 24 06:17 sample_recon_ep43_100.wav
-rw-rw-r--  1 ubuntu ubuntu 382542 Mar 24 06:17 sample_recon_ep43_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 06:07 sample_reconstruct_v3.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 07:01 sample_recon_v3b_0.wav
-rw-rw-r--  1 ubuntu ubuntu 357966 Mar 24 07:01 sample_recon_v3b_100.wav
-rw-rw-r--  1 ubuntu ubuntu 382542 Mar 24 07:01 sample_recon_v3b_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 06:34 sample_recon_v3b.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 08:09 sample_recon_v3c_0.wav
-rw-rw-r--  1 ubuntu ubuntu 358990 Mar 24 08:09 sample_recon_v3c_100.wav
-rw-rw-r--  1 ubuntu ubuntu 383566 Mar 24 08:09 sample_recon_v3c_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 14:38 sample_recon_v4_0.wav
-rw-rw-r--  1 ubuntu ubuntu 357966 Mar 24 14:38 sample_recon_v4_100.wav
-rw-rw-r--  1 ubuntu ubuntu 382542 Mar 24 14:38 sample_recon_v4_500.wav
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 10:25 samples_v3c
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 15:02 samples_v4c
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 15:10 samples_v4_ep120
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 15:44 samples_v5
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 16:04 samples_v5b
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 18:33 samples_v5_cfg
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 17:13 samples_v5_ep120
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 17:12 samples_v5_ep50
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 19:59 samples_v5_mask
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 18:08 samples_v5_newtext
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 19:26 samples_v5_pred
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 17:55 samples_v5_pred_ep44
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 18:04 samples_v5_prompted
-rw-rw-r--  1 ubuntu ubuntu 411214 Mar 24 06:17 sample_synth_ep43.wav
-rw-rw-r--  1 ubuntu ubuntu 411214 Mar 24 06:46 sample_synth_v3b.wav
-rw-rw-r--  1 ubuntu ubuntu 161358 Mar 24 06:07 sample_synth_v3.wav
-rw-rw-r--  1 ubuntu ubuntu 411214 Mar 24 14:38 sample_synth_v4_0.wav
-rw-rw-r--  1 ubuntu ubuntu 411214 Mar 24 14:38 sample_synth_v4_1.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 06:50 sample_tf_0.wav
-rw-rw-r--  1 ubuntu ubuntu 357966 Mar 24 06:50 sample_tf_100.wav
-rw-rw-r--  1 ubuntu ubuntu 382542 Mar 24 06:50 sample_tf_500.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 07:01 sample_tf_v3b_0.wav
-rw-rw-r--  1 ubuntu ubuntu 357966 Mar 24 07:01 sample_tf_v3b_100.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 08:09 sample_tf_v3c_0.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 14:38 sample_tf_v4_0.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 08:14 test_gt.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 08:14 test_noisy_same_l1.wav
-rw-rw-r--  1 ubuntu ubuntu 456782 Mar 24 07:03 test_our_mel_vocos.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 08:14 test_recon_v3c.wav
-rw-rw-r--  1 ubuntu ubuntu 456782 Mar 24 07:03 test_vocos_roundtrip.wav
-rw-rw-r--  1 ubuntu ubuntu   9539 Mar 24 14:02 train.py
-rw-rw-r--  1 ubuntu ubuntu  16092 Mar 24 07:02 train_v3b.log
-rw-rw-r--  1 ubuntu ubuntu  33886 Mar 24 09:19 train_v3c.log
-rw-rw-r--  1 ubuntu ubuntu   8401 Mar 24 06:19 train_v3.log
-rw-rw-r--  1 ubuntu ubuntu  33927 Mar 24 15:29 train_v4.log
-rw-rw-r--  1 ubuntu ubuntu  22393 Mar 24 17:17 train_v5.log
-rw-rw-r--  1 ubuntu ubuntu   5664 Mar 24 18:38 train_v5_pred.log
-rw-rw-r--  1 ubuntu ubuntu  12872 Mar 24 19:41 train_v5_pred.py
-rw-rw-r--  1 ubuntu ubuntu   7529 Mar 24 15:50 train_v5.py
drwxrwxr-x  5 ubuntu ubuntu   4096 Mar 23 22:05 v2
=== Subdirectories ===
/home/ubuntu/lewm-tts
/home/ubuntu/lewm-tts/output_v3
/home/ubuntu/lewm-tts/output_v3/logs
/home/ubuntu/lewm-tts/samples_v5
/home/ubuntu/lewm-tts/samples_v5_prompted
/home/ubuntu/lewm-tts/samples_v3c
/home/ubuntu/lewm-tts/samples_v5_newtext
/home/ubuntu/lewm-tts/output
/home/ubuntu/lewm-tts/output/logs
/home/ubuntu/lewm-tts/samples_v5b
/home/ubuntu/lewm-tts/processed_data_codec
/home/ubuntu/lewm-tts/output_v4
/home/ubuntu/lewm-tts/output_v4/logs
/home/ubuntu/lewm-tts/output_v5_pred_mask
/home/ubuntu/lewm-tts/output_v5_pred_mask/logs
/home/ubuntu/lewm-tts/output_v3c
/home/ubuntu/lewm-tts/output_v3c/logs
/home/ubuntu/lewm-tts/samples_v5_pred_ep44
/home/ubuntu/lewm-tts/samples_v5_cfg
/home/ubuntu/lewm-tts/v2


Now let me check the v2 directory which likely has the main code, and examine all the model and training files:


find /home/ubuntu/lewm-tts/v2 -type f -name "*.py" | head -30


/home/ubuntu/lewm-tts/v2/inference.py
/home/ubuntu/lewm-tts/v2/model.py
/home/ubuntu/lewm-tts/v2/preprocess.py
/home/ubuntu/lewm-tts/v2/train.py
/home/ubuntu/lewm-tts/v2/dataset.py


Perfect! Now let me systematically read the key files. I'll start with the model files to understand the architecture:


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


Excellent! Now let me read the v5 model and training files to get the complete picture:


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


Now let me check the v5 model and training files, plus the inference pipeline:


1	"""
2	LeWM TTS v5 — JEPA with EnCodec backend.
3
4	Architecture:
5	  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
6	  - Simple linear projections: 128d ↔ d_model (no lossy compression)
7	  - TextEncoder: byte-level transformer for Hindi text
8	  - DurationPredictor: predicts per-character duration
9	  - JEPAPredictor: causal transformer with frame-level text conditioning
10	"""
11
12	import math
13	import copy
14	import torch
15	import torch.nn as nn
16	import torch.nn.functional as F
17
18
19	# ─── Duration Predictor ─────────────────────────────────────────────────────
20
21	class DurationPredictor(nn.Module):
22	    """Predicts log-duration per text token. Similar to FastSpeech."""
23	    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
24	        super().__init__()
25	        layers = []
26	        for _ in range(num_layers):
27	            layers.extend([
28	                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
29	                nn.ReLU(),
30	                nn.LayerNorm(d_model),
31	                nn.Dropout(dropout),
32	            ])
33	        self.convs = nn.ModuleList(layers)
34	        self.proj = nn.Linear(d_model, 1)
35
36	    def forward(self, text_emb, text_mask=None):
37	        """
38	        text_emb: [B, T_text, d_model]
39	        Returns: log_durations [B, T_text]
40	        """
41	        x = text_emb.transpose(1, 2)  # [B, d_model, T]
42	        for i in range(0, len(self.convs), 4):
43	            x = self.convs[i](x)      # Conv1d
44	            x = self.convs[i+1](x)    # ReLU
45	            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
46	            x = self.convs[i+2](x)    # LayerNorm
47	            x = self.convs[i+3](x)    # Dropout
48	            x = x.transpose(1, 2)     # back to [B, d_model, T]
49	        x = x.transpose(1, 2)  # [B, T, d_model]
50	        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
51	        if text_mask is not None:
52	            log_dur = log_dur.masked_fill(text_mask, 0.0)
53	        return log_dur
54
55
56	def length_regulate(text_emb, durations, text_mask=None):
57	    """
58	    Expand text embeddings by repeating each token according to its duration.
59	    text_emb: [B, T_text, d_model]
60	    durations: [B, T_text] integer durations
61	    Returns: expanded [B, T_audio, d_model], where T_audio = max(sum(durations))
62	    """
63	    B, T_text, D = text_emb.shape
64	    if text_mask is not None:
65	        durations = durations.masked_fill(text_mask, 0)
66	    max_audio_len = durations.sum(dim=1).max().item()
67	    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
68	    for b in range(B):
69	        durs_b = durations[b]  # [T_text]
70	        total = durs_b.sum().item()
71	        if total == 0:
72	            continue
73	        # Use repeat_interleave — vectorized per sample
74	        exp_b = torch.repeat_interleave(text_emb[b], durs_b, dim=0)  # [total, D]
75	        expanded[b, :exp_b.shape[0]] = exp_b
76	    return expanded
77
78
79	def compute_uniform_durations(text_lengths, audio_lengths):
80	    """
81	    Compute uniform durations: spread audio frames evenly across text characters.
82	    text_lengths: [B] number of valid text tokens
83	    audio_lengths: [B] number of valid audio frames
84	    Returns: durations [B, max_text_len] integer
85	    """
86	    B = text_lengths.shape[0]
87	    max_text_len = text_lengths.max().item()
88	    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
89	    # Vectorized per sample (text lengths vary)
90	    for b in range(B):
91	        tl = text_lengths[b].item()
92	        al = audio_lengths[b].item()
93	        if tl == 0:
94	            continue
95	        base = al // tl
96	        remainder = al % tl
97	        durations[b, :tl] = base
98	        durations[b, :remainder] += 1
99	    return durations
100
101
102	# ─── JEPA Predictor ─────────────────────────────────────────────────────────
103
104	class JEPAPredictor(nn.Module):
105	    """Causal predictor with frame-level text added to audio input."""
106	    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
107	        super().__init__()
108	        self.d_model = d_model
109	        self.audio_input_proj = nn.Linear(d_model, d_model)
110	        self.text_input_proj = nn.Linear(d_model, d_model)
111	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
112	        self.dropout = nn.Dropout(dropout)
113
114	        # Use encoder (self-attention only) since text is now part of input
115	        encoder_layer = nn.TransformerEncoderLayer(
116	            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
117	            dropout=dropout, batch_first=True, activation="gelu",
118	        )
119	        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
120	        self.output_proj = nn.Sequential(
121	            nn.Linear(d_model, d_model),
122	            nn.GELU(),
123	            nn.Linear(d_model, d_model),
124	        )
125
126	    def forward(self, audio_emb, text_emb_expanded, audio_mask=None):
127	        """
128	        audio_emb: [B, T, d_model] — shifted audio embeddings
129	        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
130	        """
131	        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
132	        x = self.pos_embed(x)
133	        x = self.dropout(x)
134	        T = x.shape[1]
135	        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
136	        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
137	        return self.output_proj(predicted)
138
139	    # Keep old cross-attention interface for backward compat loading
140	    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
141	        """Legacy cross-attention forward (not used in duration-predictor mode)."""
142	        x = self.audio_input_proj(audio_emb)
143	        x = self.pos_embed(x)
144	        x = self.dropout(x)
145	        return self.output_proj(x)
146
147	    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
148	        d = mha.embed_dim
149	        nhead = mha.num_heads
150	        head_dim = d // nhead
151	        B = q_input.shape[0]
152	        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
153	        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
154	        q = F.linear(q_input, Wq, bq)
155	        k_new = F.linear(kv_input, Wk, bk)
156	        v_new = F.linear(kv_input, Wv, bv)
157	        if cache_k is not None and cache_k.shape[1] > 0:
158	            k = torch.cat([cache_k, k_new], dim=1)
159	            v = torch.cat([cache_v, v_new], dim=1)
160	        else:
161	            k, v = k_new, v_new
162	        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
163	        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
164	        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
165	        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
166	        attn = F.softmax(attn, dim=-1)
167	        out = torch.matmul(attn, v_mh)
168	        out = out.transpose(1, 2).contiguous().view(B, 1, d)
169	        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
170	        return out, k, v
171
172	    def init_cache(self, num_layers):
173	        return {
174	            'self_k': [None] * num_layers,
175	            'self_v': [None] * num_layers,
176	        }
177
178	    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
179	        """
180	        new_emb: [B, 1, d_model] — audio embedding for this step
181	        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
182	        """
183	        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
184	        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
185	        for i, layer in enumerate(self.transformer.layers):
186	            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
187	                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
188	            )
189	            x = layer.norm1(x + layer.dropout1(sa_out))
190	            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
191	            x = layer.norm2(x + layer.dropout2(ff_out))
192	        return self.output_proj(x), cache
193
194
195	# ─── Positional Encoding ────────────────────────────────────────────────────
196
197	class SinusoidalPositionalEncoding(nn.Module):
198	    def __init__(self, d_model, max_len=8192):
199	        super().__init__()
200	        pe = torch.zeros(max_len, d_model)
201	        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
202	        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
203	        pe[:, 0::2] = torch.sin(position * div_term)
204	        pe[:, 1::2] = torch.cos(position * div_term)
205	        self.register_buffer("pe", pe.unsqueeze(0))
206
207	    def forward(self, x):
208	        return x + self.pe[:, :x.shape[1]]
209
210
211	# ─── Text Encoder ───────────────────────────────────────────────────────────
212
213	class TextEncoder(nn.Module):
214	    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
215	        super().__init__()
216	        self.d_model = d_model
217	        self.char_embed = nn.Embedding(vocab_size, d_model)
218	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
219	        self.dropout = nn.Dropout(dropout)
220	        encoder_layer = nn.TransformerEncoderLayer(
221	            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
222	            dropout=dropout, batch_first=True, activation="gelu",
223	        )
224	        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
225	        self.proj = nn.Linear(d_model, d_model)
226
227	    def forward(self, text_tokens, text_mask=None):
228	        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
229	        x = self.pos_embed(x)
230	        x = self.dropout(x)
231	        x = self.transformer(x, src_key_padding_mask=text_mask)
232	        return self.proj(x)
233
234
235	# ─── Full v5 Model ──────────────────────────────────────────────────────────
236
237	class LeWMTTSv5(nn.Module):
238	    """
239	    JEPA TTS with EnCodec backend + duration predictor.
240	    Text is expanded to frame level and added to audio input — can't be ignored.
241
242	    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
243	              Duration predictor trained with uniform alignment as supervision.
244	    Inference: text → duration predict → expand text → AR predict → proj_out → decode
245	    """
246
247	    def __init__(self, config):
248	        super().__init__()
249	        d_model = config.get("d_model", 256)
250	        nhead = config.get("nhead", 4)
251	        codec_dim = config.get("codec_dim", 128)
252	        text_vocab = config.get("text_vocab_size", 256)
253	        text_layers = config.get("text_encoder_layers", 4)
254	        predictor_layers = config.get("predictor_layers", 6)
255	        dropout = config.get("dropout", 0.1)
256
257	        self.codec_dim = codec_dim
258
259	        self.text_encoder = TextEncoder(
260	            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
261	            num_layers=text_layers, dropout=dropout,
262	        )
263
264	        # Duration predictor
265	        self.duration_predictor = DurationPredictor(
266	            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
267	        )
268
269	        # Simple linear projections — no lossy compression
270	        self.proj_in = nn.Sequential(
271	            nn.Linear(codec_dim, d_model),
272	            nn.GELU(),
273	            nn.Linear(d_model, d_model),
274	        )
275	        self.proj_out = nn.Sequential(
276	            nn.Linear(d_model, d_model),
277	            nn.GELU(),
278	            nn.Linear(d_model, codec_dim),
279	        )
280
281	        self.predictor = JEPAPredictor(
282	            d_model=d_model, nhead=nhead,
283	            num_layers=predictor_layers, dropout=dropout,
284	        )
285
286	        # EMA on proj_in (provides stable prediction targets)
287	        self.ema_proj_in = copy.deepcopy(self.proj_in)
288	        for p in self.ema_proj_in.parameters():
289	            p.requires_grad = False
290	        self.ema_decay = config.get("ema_decay", 0.998)
291
292	        # Learnable start embedding
293	        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
294
295	        # Loss weights
296	        self.pred_weight = config.get("pred_weight", 10.0)
297	        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
298	        self.dur_weight = config.get("dur_weight", 1.0)
299
300	        # Input noise for train/inference gap
301	        self.input_noise = config.get("input_noise", 0.0)
302
303	        self.config = config
304
305	    @torch.no_grad()
306	    def update_ema(self):
307	        for p_online, p_ema in zip(self.proj_in.parameters(),
308	                                   self.ema_proj_in.parameters()):
309	            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
310
311	    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
312	        """
313	        Args:
314	            codec_emb: [B, 128, T] — continuous EnCodec embeddings
315	            text_tokens: [B, T_text]
316	            codec_mask: [B, T] bool, True = padding
317	            text_mask: [B, T_text] bool
318	        """
319	        B = codec_emb.shape[0]
320	        T = codec_emb.shape[2]
321
322	        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
323	        codec_seq = codec_emb.transpose(1, 2)
324
325	        # Text encoding
326	        text_emb = self.text_encoder(text_tokens, text_mask)
327
328	        # Duration prediction + ground truth durations (uniform alignment)
329	        log_dur_pred = self.duration_predictor(text_emb, text_mask)
330
331	        # Compute actual text/audio lengths from masks
332	        if text_mask is not None:
333	            text_lengths = (~text_mask).sum(dim=1)  # [B]
334	        else:
335	            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)
336	        if codec_mask is not None:
337	            audio_lengths = (~codec_mask).sum(dim=1)  # [B]
338	        else:
339	            audio_lengths = torch.full((B,), T, device=codec_emb.device)
340
341	        # Ground truth durations (uniform distribution)
342	        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
343	        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))
344
345	        # Duration loss (MSE on log-durations)
346	        if text_mask is not None:
347	            valid_text = (~text_mask).float()
348	            dur_loss = F.mse_loss(
349	                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
350	            ) / (valid_text.sum() + 1e-8)
351	        else:
352	            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)
353
354	        # Expand text to frame level using GT durations (training uses GT)
355	        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
356	        # Trim/pad to match audio length T
357	        if text_emb_expanded.shape[1] > T:
358	            text_emb_expanded = text_emb_expanded[:, :T]
359	        elif text_emb_expanded.shape[1] < T:
360	            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
361	                            device=text_emb.device)
362	            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)
363
364	        # Project codec embeddings to predictor space
365	        z = self.proj_in(codec_seq)  # [B, T, d_model]
366
367	        # EMA targets (deterministic)
368	        with torch.no_grad():
369	            self.ema_proj_in.eval()
370	            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
371	            self.ema_proj_in.train()
372
373	        # Prepend start embedding, shift right
374	        start = self.start_emb.expand(B, -1, -1)
375	        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]
376
377	        # Input noise
378	        if self.training and self.input_noise > 0:
379	            noise = torch.zeros_like(input_emb)
380	            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
381	                                        device=input_emb.device) * self.input_noise
382	            input_emb = input_emb + noise
383
384	        # Masks
385	        if codec_mask is not None:
386	            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
387	            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
388	            loss_mask = codec_mask
389	        else:
390	            pred_mask = None
391	            loss_mask = None
392
393	        # JEPA prediction with frame-level text
394	        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)
395
396	        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
397	        if loss_mask is not None:
398	            valid = (~loss_mask).unsqueeze(-1)
399	            prediction_loss = F.mse_loss(
400	                predicted * valid, target_emb * valid, reduction="sum"
401	            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
402	        else:
403	            prediction_loss = F.mse_loss(predicted, target_emb)
404
405	        # ─── Loss 2: Roundtrip (proj_in → proj_out → match original codec_emb) ───
406	        codec_recon = self.proj_out(z)  # [B, T, 128]
407
408	        if codec_mask is not None:
409	            valid_rt = (~codec_mask).unsqueeze(-1).float()
410	            roundtrip_loss = F.l1_loss(
411	                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
412	            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
413	        else:
414	            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)
415
416	        total_loss = (self.pred_weight * prediction_loss
417	                     + self.roundtrip_weight * roundtrip_loss
418	                     + self.dur_weight * dur_loss)
419
420	        return {
421	            "total_loss": total_loss,
422	            "prediction_loss": prediction_loss,
423	            "roundtrip_loss": roundtrip_loss,
424	            "dur_loss": dur_loss,
425	        }
426
427	    def predict_durations(self, text_emb, text_mask=None):
428	        """Predict durations from text embeddings (for inference)."""
429	        log_dur = self.duration_predictor(text_emb, text_mask)
430	        durations = torch.exp(log_dur).round().long().clamp(min=1)
431	        if text_mask is not None:
432	            durations = durations.masked_fill(text_mask, 0)
433	        return durations
434
435	    def init_ar_cache(self):
436	        num_layers = len(self.predictor.transformer.layers)
437	        return self.predictor.init_cache(num_layers)
438
439	    def predict_next_cached(self, new_emb, text_emb_frame, step_idx, cache):
440	        """
441	        new_emb: [B, 1, d_model]
442	        text_emb_frame: [B, 1, d_model] — expanded text for this frame
443	        """
444	        return self.predictor.inference_step(new_emb, text_emb_frame, step_idx, cache)
445
446
447	def count_parameters(model):
448	    return sum(p.numel() for p in model.parameters() if p.requires_grad)
449
450
451	def build_model(config=None):
452	    if config is None:
453	        config = {
454	            "d_model": 256, "nhead": 4, "codec_dim": 128,
455	            "text_vocab_size": 256, "text_encoder_layers": 4,
456	            "predictor_layers": 6, "dropout": 0.1,
457	            "pred_weight": 10.0, "roundtrip_weight": 1.0,
458	            "ema_decay": 0.998, "input_noise": 0.0,
459	        }
460	    model = LeWMTTSv5(config)
461	    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
462	    return model, config
463


Now let me check the inference 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	        print(f"Model loaded from {checkpoint_path}")
37
38	    def vocos_decode(self, mel):
39	        """Decode mel [B, 100, T] → waveform via Vocos."""
40	        with torch.no_grad():
41	            waveform = self.vocos.decode(mel)  # [B, T_audio]
42	        return waveform
43
44	    def text_to_tokens(self, text):
45	        tokens = list(text.encode("utf-8"))
46	        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
47
48	    @torch.no_grad()
49	    def synthesize(self, text, max_steps=300, temperature=0.0):
50	        """Generate speech from text using AR prediction with KV-cache."""
51	        text_tokens = self.text_to_tokens(text)
52	        text_emb = self.model.text_encoder(text_tokens)
53
54	        # Use learned start embedding instead of random noise
55	        start_emb = self.model.start_emb.to(self.device)
56	        all_embeddings = [start_emb]
57
58	        cache = self.model.init_ar_cache()
59	        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
60	        if temperature > 0:
61	            next_emb = next_emb + torch.randn_like(next_emb) * temperature
62	        all_embeddings.append(next_emb)
63
64	        prev_norms = [next_emb.norm().item()]
65
66	        for step in range(1, max_steps):
67	            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
68	            if temperature > 0:
69	                next_emb = next_emb + torch.randn_like(next_emb) * temperature
70	            all_embeddings.append(next_emb)
71
72	            norm = next_emb.norm().item()
73	            prev_norms.append(norm)
74	            if len(prev_norms) > 20:
75	                recent = prev_norms[-20:]
76	                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
77	                    break
78
79	        audio_embs = torch.cat(all_embeddings, dim=1)
80	        mel = self.model.mel_decoder(audio_embs)
81	        waveform = self.vocos_decode(mel.cpu())
82	        return waveform.squeeze().numpy(), 24000
83
84	    @torch.no_grad()
85	    def reconstruct(self, mel_path):
86	        """Teacher-forced reconstruction: encode real mel → decode back."""
87	        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
88
89	        z, mu, logvar = self.model.audio_encoder(mel)
90	        mel_recon = self.model.mel_decoder(mu)
91
92	        waveform = self.vocos_decode(mel_recon.cpu())
93	        return waveform.squeeze().numpy(), 24000
94
95	    @torch.no_grad()
96	    def reconstruct_with_prediction(self, mel_path):
97	        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
98	        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
99
100	        z, mu, logvar = self.model.audio_encoder(mel)
101
102	        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
103	        text_emb = self.model.text_encoder(text_tokens)
104
105	        predicted = self.model.predictor(mu[:, :-1], text_emb)
106	        combined = torch.cat([mu[:, :1], predicted], dim=1)
107
108	        mel_recon = self.model.mel_decoder(combined)
109	        waveform = self.vocos_decode(mel_recon.cpu())
110	        return waveform.squeeze().numpy(), 24000
111
112	    def save_audio(self, waveform, sr, output_path):
113	        # Normalize
114	        mx = np.abs(waveform).max()
115	        if mx > 0:
116	            waveform = waveform / mx * 0.95
117	        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
118	        torchaudio.save(output_path, wav_tensor, sr)
119	        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
120
121
122	def main():
123	    parser = argparse.ArgumentParser()
124	    parser.add_argument("--checkpoint", required=True)
125	    parser.add_argument("--text", default=None, help="Text to synthesize")
126	    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
127	    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
128	                        default="synthesize")
129	    parser.add_argument("--output", default="output.wav")
130	    parser.add_argument("--max_steps", type=int, default=300)
131	    parser.add_argument("--temperature", type=float, default=0.0)
132	    args = parser.parse_args()
133
134	    tts = LeWMTTSInference(args.checkpoint)
135
136	    if args.mode == "synthesize":
137	        assert args.text, "Need --text for synthesis"
138	        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
139	    elif args.mode == "reconstruct":
140	        assert args.mel, "Need --mel for reconstruction"
141	        waveform, sr = tts.reconstruct(args.mel)
142	    elif args.mode == "predict":
143	        assert args.mel, "Need --mel for prediction test"
144	        waveform, sr = tts.reconstruct_with_prediction(args.mel)
145
146	    tts.save_audio(waveform, sr, args.output)
147
148
149	if __name__ == "__main__":
150	    main()
151


Now let me check the v5 inference and training files:


1	"""
2	Inference for LeWM TTS v5 (codec-based JEPA with duration predictor).
3	Text → duration predict → expand text → JEPA predict → proj_out → EnCodec decode → 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_v5 import LeWMTTSv5, length_regulate
13
14
15	class LeWMTTSv5Inference:
16	    def __init__(self, checkpoint_path, device="cuda"):
17	        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
18
19	        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
20	        config = ckpt["config"]
21
22	        self.model = LeWMTTSv5(config).to(self.device)
23	        self.model.load_state_dict(ckpt["model"])
24	        self.model.eval()
25
26	        # EnCodec decoder (frozen)
27	        from encodec import EncodecModel
28	        self.codec = EncodecModel.encodec_model_24khz()
29	        self.codec.set_target_bandwidth(6.0)
30	        self.codec.eval().to(self.device)
31
32	        self.config = config
33	        print(f"v5 model loaded from {checkpoint_path}")
34
35	    def text_to_tokens(self, text):
36	        tokens = list(text.encode("utf-8"))
37	        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
38
39	    @torch.no_grad()
40	    def synthesize(self, text, duration_scale=1.0, temperature=0.0):
41	        """AR synthesis with duration predictor for text alignment."""
42	        text_tokens = self.text_to_tokens(text)
43	        text_emb = self.model.text_encoder(text_tokens)
44
45	        # Predict durations and expand text
46	        durations = self.model.predict_durations(text_emb)
47	        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
48	        total_frames = durations.sum().item()
49	        print(f"Text: {len(text)} chars → {total_frames} frames ({total_frames/75:.2f}s)")
50
51	        # Expand text to frame level
52	        text_emb_expanded = length_regulate(text_emb, durations)  # [1, T_total, d_model]
53
54	        # AR generation
55	        start_emb = self.model.start_emb.to(self.device)
56	        cache = self.model.init_ar_cache()
57	        all_embeddings = [start_emb]
58
59	        # First step
60	        text_frame = text_emb_expanded[:, 0:1]  # [1, 1, d_model]
61	        next_emb, cache = self.model.predict_next_cached(start_emb, text_frame, 0, cache)
62	        if temperature > 0:
63	            next_emb = next_emb + torch.randn_like(next_emb) * temperature
64	        all_embeddings.append(next_emb)
65
66	        for step in range(1, total_frames):
67	            # Get text embedding for this frame
68	            if step < text_emb_expanded.shape[1]:
69	                text_frame = text_emb_expanded[:, step:step+1]
70	            else:
71	                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)
72
73	            next_emb, cache = self.model.predict_next_cached(next_emb, text_frame, step, cache)
74	            if temperature > 0:
75	                next_emb = next_emb + torch.randn_like(next_emb) * temperature
76	            all_embeddings.append(next_emb)
77
78	        # Concat predicted embeddings and project to codec space
79	        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
80	        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
81	        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]
82
83	        # EnCodec decode
84	        waveform = self.codec.decoder(codec_emb)
85	        waveform = waveform.squeeze().cpu().numpy()
86
87	        return waveform, 24000
88
89	    @torch.no_grad()
90	    def synthesize_prompted(self, text, prompt_audio_path, prompt_text=None,
91	                            duration_scale=1.0, temperature=0.0):
92	        """AR synthesis with audio prompt as prefix for voice/style reference."""
93	        from encodec.utils import convert_audio
94
95	        # Encode prompt audio
96	        wav, sr = torchaudio.load(prompt_audio_path)
97	        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
98	        prompt_codec = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T_prompt]
99	        T_prompt = prompt_codec.shape[2]
100	        prompt_z = self.model.proj_in(prompt_codec.transpose(1, 2))  # [1, T_prompt, d_model]
101
102	        # Text encoding (full text = prompt_text + synthesis text)
103	        if prompt_text:
104	            full_text = prompt_text + text
105	        else:
106	            full_text = text
107	        text_tokens = self.text_to_tokens(full_text)
108	        text_emb = self.model.text_encoder(text_tokens)
109
110	        # Predict durations and expand text to frame level
111	        durations = self.model.predict_durations(text_emb)
112	        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
113	        total_frames = durations.sum().item()
114	        text_emb_expanded = length_regulate(text_emb, durations)
115
116	        # How many new frames to generate beyond the prompt
117	        gen_frames = max(total_frames - T_prompt, total_frames // 2)
118	        print(f"Prompt: {T_prompt} frames ({T_prompt/75:.2f}s) | "
119	              f"Generate: {gen_frames} frames ({gen_frames/75:.2f}s) | "
120	              f"Total text frames: {total_frames}")
121
122	        # Feed prompt through predictor to build KV cache
123	        cache = self.model.init_ar_cache()
124
125	        # Process prompt: start_emb + prompt_z[:-1] as input
126	        start_emb = self.model.start_emb.to(self.device)
127	        prompt_input = torch.cat([start_emb, prompt_z[:, :-1]], dim=1)  # [1, T_prompt, d_model]
128
129	        # Feed prompt frames one by one to build cache
130	        for step in range(T_prompt):
131	            frame_emb = prompt_input[:, step:step+1]
132	            if step < text_emb_expanded.shape[1]:
133	                text_frame = text_emb_expanded[:, step:step+1]
134	            else:
135	                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)
136	            next_emb, cache = self.model.predict_next_cached(frame_emb, text_frame, step, cache)
137
138	        # AR generation continuing from prompt
139	        all_embeddings = [prompt_z]  # keep prompt embeddings
140	        cur_emb = prompt_z[:, -1:]   # last prompt frame as input
141
142	        for step in range(gen_frames):
143	            abs_step = T_prompt + step
144	            if abs_step < text_emb_expanded.shape[1]:
145	                text_frame = text_emb_expanded[:, abs_step:abs_step+1]
146	            else:
147	                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)
148
149	            next_emb, cache = self.model.predict_next_cached(cur_emb, text_frame, abs_step, cache)
150	            if temperature > 0:
151	                next_emb = next_emb + torch.randn_like(next_emb) * temperature
152	            all_embeddings.append(next_emb)
153	            cur_emb = next_emb
154
155	        # Decode
156	        pred_embs = torch.cat(all_embeddings, dim=1)
157	        codec_emb = self.model.proj_out(pred_embs).transpose(1, 2)
158	        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
159
160	        return waveform, 24000
161
162	    @torch.no_grad()
163	    def reconstruct(self, audio_path):
164	        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
165	        from encodec.utils import convert_audio
166	        wav, sr = torchaudio.load(audio_path)
167	        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
168
169	        codec_emb = self.codec.encoder(wav.unsqueeze(0))
170	        z = self.model.proj_in(codec_emb.transpose(1, 2))
171	        codec_recon = self.model.proj_out(z).transpose(1, 2)
172	        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
173	        return waveform, 24000
174
175	    @torch.no_grad()
176	    def reconstruct_codec_only(self, audio_path):
177	        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
178	        from encodec.utils import convert_audio
179	        wav, sr = torchaudio.load(audio_path)
180	        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
181
182	        codec_emb = self.codec.encoder(wav.unsqueeze(0))
183	        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
184	        return waveform, 24000
185
186	    @torch.no_grad()
187	    def reconstruct_with_prediction(self, audio_path, text):
188	        """Encode audio → proj_in → teacher-forced predict with text alignment → proj_out → decode."""
189	        from encodec.utils import convert_audio
190	        from model_v5 import compute_uniform_durations
191	        wav, sr = torchaudio.load(audio_path)
192	        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
193
194	        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
195	        T = codec_emb.shape[2]
196	        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]
197
198	        text_tokens = self.text_to_tokens(text)
199	        text_emb = self.model.text_encoder(text_tokens)
200
201	        # Compute uniform durations and expand text
202	        text_len = torch.tensor([text_tokens.shape[1]], device=self.device)
203	        audio_len = torch.tensor([T], device=self.device)
204	        durations = compute_uniform_durations(text_len, audio_len)
205	        text_emb_expanded = length_regulate(text_emb, durations)
206	        if text_emb_expanded.shape[1] > T:
207	            text_emb_expanded = text_emb_expanded[:, :T]
208	        elif text_emb_expanded.shape[1] < T:
209	            pad = torch.zeros(1, T - text_emb_expanded.shape[1], text_emb.shape[-1],
210	                            device=self.device)
211	            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)
212
213	        B = z.shape[0]
214	        start = self.model.start_emb.expand(B, -1, -1)
215	        input_emb = torch.cat([start, z[:, :-1]], dim=1)
216	        predicted = self.model.predictor(input_emb, text_emb_expanded)
217
218	        codec_recon = self.model.proj_out(predicted).transpose(1, 2)
219	        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
220	        return waveform, 24000
221
222	    def save_audio(self, waveform, sr, output_path):
223	        mx = np.abs(waveform).max()
224	        if mx > 0:
225	            waveform = waveform / mx * 0.95
226	        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
227	        torchaudio.save(output_path, wav_tensor, sr)
228	        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
229
230
231	def main():
232	    parser = argparse.ArgumentParser()
233	    parser.add_argument("--checkpoint", required=True)
234	    parser.add_argument("--text", default=None)
235	    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
236	    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
237	                        default="synthesize")
238	    parser.add_argument("--output", default="output_v5.wav")
239	    parser.add_argument("--duration_scale", type=float, default=1.0)
240	    parser.add_argument("--temperature", type=float, default=0.0)
241	    args = parser.parse_args()
242
243	    tts = LeWMTTSv5Inference(args.checkpoint)
244
245	    if args.mode == "synthesize":
246	        assert args.text, "Need --text"
247	        waveform, sr = tts.synthesize(args.text, args.duration_scale, args.temperature)
248	    elif args.mode == "reconstruct":
249	        assert args.audio, "Need --audio"
250	        waveform, sr = tts.reconstruct(args.audio)
251	    elif args.mode == "codec_only":
252	        assert args.audio, "Need --audio"
253	        waveform, sr = tts.reconstruct_codec_only(args.audio)
254	    elif args.mode == "predict":
255	        assert args.audio and args.text, "Need --audio and --text"
256	        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)
257
258	    tts.save_audio(waveform, sr, args.output)
259
260
261	if __name__ == "__main__":
262	    main()
263


Now let me check the training files and dataset definitions:


1	"""
2	Training loop for LeWM TTS v5 (codec-based JEPA).
3	"""
4
5	import os
6	import json
7	import time
8	import argparse
9	import torch
10	from torch.utils.tensorboard import SummaryWriter
11	from pathlib import Path
12
13	from model_v5 import build_model
14	from dataset_v5 import build_dataloader
15
16
17	def train(args):
18	    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
19	    print(f"Device: {device}")
20
21	    config = {
22	        "d_model": args.d_model,
23	        "nhead": args.nhead,
24	        "codec_dim": 128,
25	        "text_vocab_size": 256,
26	        "text_encoder_layers": args.text_layers,
27	        "predictor_layers": args.predictor_layers,
28	        "dropout": args.dropout,
29	        "pred_weight": args.pred_weight,
30	        "roundtrip_weight": args.roundtrip_weight,
31	        "ema_decay": args.ema_decay,
32	        "input_noise": 0.0,
33	    }
34	    model, config = build_model(config)
35	    model = model.to(device)
36
37	    manifest_path = os.path.join(args.data_dir, "manifest.json")
38	    loader, dataset = build_dataloader(
39	        manifest_path, batch_size=args.batch_size,
40	        num_workers=args.num_workers, max_codec_frames=args.max_codec_frames,
41	    )
42	    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")
43
44	    optimizer = torch.optim.AdamW(
45	        model.parameters(), lr=args.lr,
46	        betas=(0.9, 0.98), weight_decay=0.01,
47	    )
48
49	    total_steps = args.epochs * len(loader)
50	    warmup_steps = min(4000, total_steps // 10)
51
52	    def lr_lambda(step):
53	        if step < warmup_steps:
54	            return step / max(1, warmup_steps)
55	        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
56	        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())
57
58	    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
59
60	    ckpt_dir = Path(args.output_dir) / "checkpoints"
61	    ckpt_dir.mkdir(parents=True, exist_ok=True)
62	    writer = SummaryWriter(Path(args.output_dir) / "logs")
63
64	    with open(Path(args.output_dir) / "config.json", "w") as f:
65	        json.dump(config, f, indent=2)
66
67	    start_epoch = 0
68	    global_step = 0
69	    if args.resume:
70	        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
71	        model.load_state_dict(ckpt["model"])
72	        optimizer.load_state_dict(ckpt["optimizer"])
73	        scheduler.load_state_dict(ckpt["scheduler"])
74	        start_epoch = ckpt["epoch"] + 1
75	        global_step = ckpt["global_step"]
76	        print(f"Resumed from epoch {start_epoch}, step {global_step}")
77
78	    model.train()
79	    best_loss = float("inf")
80
81	    for epoch in range(start_epoch, args.epochs):
82	        epoch_loss = 0.0
83	        epoch_pred = 0.0
84	        epoch_rt = 0.0
85	        t0 = time.time()
86
87	        for batch_idx, batch in enumerate(loader):
88	            codec_emb = batch["codec_emb"].to(device)
89	            text_tokens = batch["text_tokens"].to(device)
90	            codec_mask = batch["codec_mask"].to(device)
91	            text_mask = batch["text_mask"].to(device)
92
93	            # Input noise annealing
94	            noise_progress = min(1.0, global_step / (total_steps * 0.3))
95	            model.input_noise = args.input_noise_max * noise_progress
96
97	            losses = model(codec_emb, text_tokens, codec_mask, text_mask)
98
99	            total_loss = losses["total_loss"]
100	            pred_loss = losses["prediction_loss"]
101	            rt_loss = losses["roundtrip_loss"]
102
103	            optimizer.zero_grad()
104	            total_loss.backward()
105	            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
106	            optimizer.step()
107	            scheduler.step()
108	            model.update_ema()
109
110	            epoch_loss += total_loss.item()
111	            epoch_pred += pred_loss.item()
112	            epoch_rt += rt_loss.item()
113	            global_step += 1
114
115	            if global_step % args.log_every == 0:
116	                lr = optimizer.param_groups[0]["lr"]
117	                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
118	                writer.add_scalar("train/prediction_loss", pred_loss.item(), global_step)
119	                writer.add_scalar("train/roundtrip_loss", rt_loss.item(), global_step)
120	                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
121	                writer.add_scalar("train/lr", lr, global_step)
122
123	            if global_step % args.print_every == 0:
124	                elapsed = time.time() - t0
125	                print(
126	                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
127	                    f"loss={total_loss.item():.4f} pred={pred_loss.item():.4f} "
128	                    f"rt={rt_loss.item():.4f} | "
129	                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"
130	                )
131
132	        n_batches = len(loader)
133	        avg_loss = epoch_loss / n_batches
134	        avg_pred = epoch_pred / n_batches
135	        avg_rt = epoch_rt / n_batches
136	        epoch_time = time.time() - t0
137
138	        print(
139	            f"Epoch {epoch+1}/{args.epochs} | "
140	            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} rt={avg_rt:.4f} | {epoch_time:.1f}s"
141	        )
142
143	        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
144	            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
145	            torch.save({
146	                "epoch": epoch,
147	                "global_step": global_step,
148	                "model": model.state_dict(),
149	                "optimizer": optimizer.state_dict(),
150	                "scheduler": scheduler.state_dict(),


Let me check the dataset files:


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


Now let me check the dataset_v5 file: