The file /home/ubuntu/lewm-tts/inference_v5.py has been updated successfully.

        return waveform, 24000

    @torch.no_grad()
    def synthesize_prompted(self, text, prompt_audio_path, prompt_text=None,
                            duration_scale=1.0, temperature=0.0):
        """AR synthesis with audio prompt as prefix for voice/style reference."""
        from encodec.utils import convert_audio

        # Encode prompt audio
        wav, sr = torchaudio.load(prompt_audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
        prompt_codec = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T_prompt]
        T_prompt = prompt_codec.shape[2]
        prompt_z = self.model.proj_in(prompt_codec.transpose(1, 2))  # [1, T_prompt, d_model]

        # Text encoding (full text = prompt_text + synthesis text)
        if prompt_text:
            full_text = prompt_text + text
        else:
            full_text = text
        text_tokens = self.text_to_tokens(full_text)
        text_emb = self.model.text_encoder(text_tokens)

        # Predict durations and expand text to frame level
        durations = self.model.predict_durations(text_emb)
        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
        total_frames = durations.sum().item()
        text_emb_expanded = length_regulate(text_emb, durations)

        # How many new frames to generate beyond the prompt
        gen_frames = max(total_frames - T_prompt, total_frames // 2)
        print(f"Prompt: {T_prompt} frames ({T_prompt/75:.2f}s) | "
              f"Generate: {gen_frames} frames ({gen_frames/75:.2f}s) | "
              f"Total text frames: {total_frames}")

        # Feed prompt through predictor to build KV cache
        cache = self.model.init_ar_cache()

        # Process prompt: start_emb + prompt_z[:-1] as input
        start_emb = self.model.start_emb.to(self.device)
        prompt_input = torch.cat([start_emb, prompt_z[:, :-1]], dim=1)  # [1, T_prompt, d_model]

        # Feed prompt frames one by one to build cache
        for step in range(T_prompt):
            frame_emb = prompt_input[:, step:step+1]
            if step < text_emb_expanded.shape[1]:
                text_frame = text_emb_expanded[:, step:step+1]
            else:
                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)
            next_emb, cache = self.model.predict_next_cached(frame_emb, text_frame, step, cache)

        # AR generation continuing from prompt
        all_embeddings = [prompt_z]  # keep prompt embeddings
        cur_emb = prompt_z[:, -1:]   # last prompt frame as input

        for step in range(gen_frames):
            abs_step = T_prompt + step
            if abs_step < text_emb_expanded.shape[1]:
                text_frame = text_emb_expanded[:, abs_step:abs_step+1]
            else:
                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

            next_emb, cache = self.model.predict_next_cached(cur_emb, text_frame, abs_step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)
            cur_emb = next_emb

        # Decode
        pred_embs = torch.cat(all_embeddings, dim=1)
        codec_emb = self.model.proj_out(pred_embs).transpose(1, 2)
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):

"""
Inference for LeWM TTS v5 (codec-based JEPA with duration predictor).
Text → duration predict → expand text → JEPA predict → proj_out → EnCodec decode → waveform
"""

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

from model_v5 import LeWMTTSv5, length_regulate

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

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

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

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

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

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

    @torch.no_grad()
    def synthesize(self, text, duration_scale=1.0, temperature=0.0):
        """AR synthesis with duration predictor for text alignment."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Predict durations and expand text
        durations = self.model.predict_durations(text_emb)
        durations = (durations.float() * duration_scale).round().long().clamp(min=1)
        total_frames = durations.sum().item()
        print(f"Text: {len(text)} chars → {total_frames} frames ({total_frames/75:.2f}s)")

        # Expand text to frame level
        text_emb_expanded = length_regulate(text_emb, durations)  # [1, T_total, d_model]

        # AR generation
        start_emb = self.model.start_emb.to(self.device)
        cache = self.model.init_ar_cache()
        all_embeddings = [start_emb]

        # First step
        text_frame = text_emb_expanded[:, 0:1]  # [1, 1, d_model]
        next_emb, cache = self.model.predict_next_cached(start_emb, text_frame, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

        for step in range(1, total_frames):
            # Get text embedding for this frame
            if step < text_emb_expanded.shape[1]:
                text_frame = text_emb_expanded[:, step:step+1]
            else:
                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

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

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z = self.model.proj_in(codec_emb.transpose(1, 2))
        codec_recon = self.model.proj_out(z).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

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

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

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

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        T = codec_emb.shape[2]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

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

        # Compute uniform durations and expand text
        text_len = torch.tensor([text_tokens.shape[1]], device=self.device)
        audio_len = torch.tensor([T], device=self.device)
        durations = compute_uniform_durations(text_len, audio_len)
        text_emb_expanded = length_regulate(text_emb, durations)
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(1, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=self.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

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

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

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

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

    tts = LeWMTTSv5Inference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

+    def synthesize_prompted(self, text, prompt_audio_path, prompt_text=None,

+                            duration_scale=1.0, temperature=0.0):

+        """AR synthesis with audio prompt as prefix for voice/style reference."""

+        from encodec.utils import convert_audio

+        wav, sr = torchaudio.load(prompt_audio_path)

+        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

+        prompt_codec = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T_prompt]

+        T_prompt = prompt_codec.shape[2]

+        prompt_z = self.model.proj_in(prompt_codec.transpose(1, 2))  # [1, T_prompt, d_model]

+        # Text encoding (full text = prompt_text + synthesis text)

+            full_text = prompt_text + text

+        text_tokens = self.text_to_tokens(full_text)

+        text_emb = self.model.text_encoder(text_tokens)

+        # Predict durations and expand text to frame level

+        durations = self.model.predict_durations(text_emb)

+        durations = (durations.float() * duration_scale).round().long().clamp(min=1)

+        total_frames = durations.sum().item()

+        text_emb_expanded = length_regulate(text_emb, durations)

+        # How many new frames to generate beyond the prompt

+        gen_frames = max(total_frames - T_prompt, total_frames // 2)

+        print(f"Prompt: {T_prompt} frames ({T_prompt/75:.2f}s) | "

+              f"Generate: {gen_frames} frames ({gen_frames/75:.2f}s) | "

+              f"Total text frames: {total_frames}")

+        # Feed prompt through predictor to build KV cache

+        cache = self.model.init_ar_cache()

+        # Process prompt: start_emb + prompt_z[:-1] as input

+        start_emb = self.model.start_emb.to(self.device)

+        prompt_input = torch.cat([start_emb, prompt_z[:, :-1]], dim=1)  # [1, T_prompt, d_model]

+        # Feed prompt frames one by one to build cache

+            frame_emb = prompt_input[:, step:step+1]

+            if step < text_emb_expanded.shape[1]:

+                text_frame = text_emb_expanded[:, step:step+1]

+                text_frame = torch.zeros(1, 1, text_emb.shape[-1], device=self.device)

+            next_emb, cache = self.model.predict_next_cached(frame_emb, text_frame, step, cache)

+        # AR generation continuing from prompt

+        all_embeddings = [prompt_z]  # keep prompt embeddings

+        cur_emb = prompt_z[:, -1:]   # last prompt frame as input

+            if abs_step < text_emb_expanded.shape[1]:

+                text_frame = text_emb_expanded[:, abs_step:abs_step+1]

+            next_emb, cache = self.model.predict_next_cached(cur_emb, text_frame, abs_step, cache)

+                next_emb = next_emb + torch.randn_like(next_emb) * temperature

+            all_embeddings.append(next_emb)

+        pred_embs = torch.cat(all_embeddings, dim=1)

+        codec_emb = self.model.proj_out(pred_embs).transpose(1, 2)

+        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()


Now let me generate samples with audio prompts:


Bash

python3 -c "
import json
with open('processed_data_100mel/manifest.json') as f:
    manifest = json.load(f)
# Pick a few short clips as prompts (3-5s)
for entry in manifest[:20]:
    if 3.0 < entry['duration'] < 6.0:
        print(f\"{entry['duration']:.1f}s | {entry['audio_path']} | {entry['text'][:80]}\")
" 2>&1 | head -5

Find good prompt clips


3.4s | /home/ubuntu/lewm-tts/processed_data_100mel/audio/000012.wav | neration के अनुरूप ढाल रहे है
3.4s | /home/ubuntu/lewm-tts/processed_data_100mel/audio/000019.wav | से जानने में world leaders ने ब


python3 -c "
import torch, torchaudio, numpy as np
from inference_v5 import LeWMTTSv5Inference

tts = LeWMTTSv5Inference('output_v5_pred_dur/checkpoints/best.pt')

prompt_audio = 'processed_data_100mel/audio/000000.wav'
prompt_text = 'मेरे प्यारे देशवासियों, नमस्कार। मन की बात में आपका स्वागत है, अभिन'

# Test 1: prompted synthesis — continue speaking with same voice
print('\\n=== Prompted synthesis ===')
wav, sr = tts.synthesize_prompted(
    text='भारत एक महान देश है और हम सब भारतवासी हैं।',
    prompt_audio_path=prompt_audio,
    prompt_text=prompt_text,
    duration_scale=1.0,
    temperature=0.0
)
tts.save_audio(wav, sr, 'samples_v5_pred/prompted_0.wav')

# Test 2: different text
wav, sr = tts.synthesize_prompted(
    text='आज हम बात करेंगे शिक्षा के बारे में।',
    prompt_audio_path=prompt_audio,
    prompt_text=prompt_text,
    duration_scale=1.0,
    temperature=0.0
)
tts.save_audio(wav, sr, 'samples_v5_pred/prompted_1.wav')

# Test 3: with slight temperature
wav, sr = tts.synthesize_prompted(
    text='भारत एक महान देश है और हम सब भारतवासी हैं।',
    prompt_audio_path=prompt_audio,
    prompt_text=prompt_text,
    duration_scale=1.0,
    temperature=0.05
)
tts.save_audio(wav, sr, 'samples_v5_pred/prompted_temp_0.wav')

# Test 4: shorter prompt
wav, sr = tts.synthesize_prompted(
    text='नमस्कार दोस्तों, कैसे हैं आप सब।',
    prompt_audio_path='processed_data_100mel/audio/000012.wav',
    prompt_text='neration के अनुरूप ढाल रहे है',
    duration_scale=1.0,
    temperature=0.0
)
tts.save_audio(wav, sr, 'samples_v5_pred/prompted_2.wav')

print('\\nDone!')
" 2>&1

Generate voice-prompted synthesis samples


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
v5 model loaded from output_v5_pred_dur/checkpoints/best.pt

=== Prompted synthesis ===
Prompt: 714 frames (9.52s) | Generate: 327 frames (4.36s) | Total text frames: 655
Saved: samples_v5_pred/prompted_0.wav (13.88s)
Prompt: 714 frames (9.52s) | Generate: 322 frames (4.29s) | Total text frames: 644
Saved: samples_v5_pred/prompted_1.wav (13.81s)
Prompt: 714 frames (9.52s) | Generate: 327 frames (4.36s) | Total text frames: 655
Saved: samples_v5_pred/prompted_temp_0.wav (13.88s)
Prompt: 252 frames (3.36s) | Generate: 242 frames (3.23s) | Total text frames: 485
Saved: samples_v5_pred/prompted_2.wav (6.59s)

Done!


Generated 4 voice-prompted samples in `samples_v5_pred/`:

| File | Prompt | Generated text |
|------|--------|----------------|
| `prompted_0.wav` | 000000.wav (9.5s) | "भारत एक महान देश है..." |
| `prompted_1.wav` | 000000.wav (9.5s) | "आज हम बात करेंगे शिक्षा के बारे में।" |
| `prompted_temp_0.wav` | same, temp=0.05 | same as 0, with slight randomness |
| `prompted_2.wav` | 000012.wav (3.4s) | "नमस्कार दोस्तों, कैसे हैं आप सब।" |

The prompt audio plays first (real audio via roundtrip), then the AR-generated continuation follows. This should tell us if the predictor can maintain voice quality when given a real audio prefix. Have a listen!


prompted_0.wav      │ 000000.wav (9.5s) │त"भारत एक महान देश है..."│             │
  ├─────────────────────┼───────────────────┼───────────────────────────────────┤
  │ prompted_1.wav      │ 000000.wav (9.5s) │ "आज हम बात करेंगे शिक्षारेके बारे में।"        │
  ├─────────────────────┼───────────────────┼───────────────────────────────────┤
  │ prompted_temp_0.wav │ same, temp=0.05   │ same as 0, with slight randomness │ for these 3 it said same text..   │ prompted_2.wav      │ 000012.wav (3.4s) │ "नमस्कार दोस्तों, कैसे हैं आप सब।"          │
  └─────────────────────┴───────────────────┴───────────────────────────────── thus si not mentioned text


The user is saying that all the prompted samples just repeat the same prompt text/audio rather than generating the new text. The AR continuation is just repeating what was in the prompt, not speaking the new text.

This makes sense - the model is trained with teacher forcing, so when it AR generates, it just continues the pattern it saw in the prompt. The predictor is collapsing / not using the text conditioning properly.

The issue is that the generated portion sounds like the prompt text, not the target text. This means the text conditioning isn't working - the predictor is ignoring the text embeddings and just continuing the audio pattern.

This is the same AR collapse issue noted in the memory: "AR synthesis collapses — predictor pred_loss=0.02 (TF) but AR free-running converges to silence. Needs scheduled sampling."

So the generated audio is repeating the prompt content instead of speaking the new text - either the text conditioning is too weak and the predictor learned to mostly copy previous audio, or the AR generation is just looping. Let me trace through the code to understand what's happening: we cache the prompt frames, then continue AR generation, and the text embeddings should cover both the prompt and new text portions, so at generation time the embeddings should correspond to the new text. But if the predictor isn't using those embeddings properly, it would just keep repeating the audio pattern it learned.

The prediction loss of 0.42 is still quite high, which suggests the predictor isn't learning well, but the real issue is likely the text conditioning itself. I should test whether text conditioning even works by doing a teacher-forced continuation - feed real audio for the prompt portion but then continue with teacher-forced prediction using the new text to see if the text encoder is actually being used at all.


So the AR continuation just repeats the prompt audio instead of speaking the new text. That means the **predictor is ignoring text conditioning** — it's learned to copy the previous frame pattern via self-attention and barely uses the text cross-attention.

This is the same core issue from before: pred_loss=0.42 in teacher-forcing mode, but AR free-running doesn't follow text.

Let me run a quick diagnostic to confirm text conditioning is weak:


python3 -c "
import torch
from inference_v5 import LeWMTTSv5Inference
from model_v5 import length_regulate, compute_uniform_durations
from encodec.utils import convert_audio
import torchaudio

tts = LeWMTTSv5Inference('output_v5_pred_dur/checkpoints/best.pt')
model = tts.model

# Load a real audio clip
wav, sr = torchaudio.load('processed_data_100mel/audio/000000.wav')
wav = convert_audio(wav, sr, 24000, 1).to(tts.device)
codec_emb = tts.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
T = codec_emb.shape[2]
z = model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

# Teacher-forced prediction with CORRECT text
correct_text = 'मेरे प्यारे देशवासियों, नमस्कार। मन की बात में आपका स्वागत है, अभिन'
correct_tokens = tts.text_to_tokens(correct_text)
correct_emb = model.text_encoder(correct_tokens)
dur_c = compute_uniform_durations(
    torch.tensor([correct_tokens.shape[1]], device=tts.device),
    torch.tensor([T], device=tts.device))
correct_expanded = length_regulate(correct_emb, dur_c)[:, :T]

# Teacher-forced prediction with WRONG text (totally different)
wrong_text = 'abcdefghijklmnopqrstuvwxyz hello world testing random english text here'
wrong_tokens = tts.text_to_tokens(wrong_text)
wrong_emb = model.text_encoder(wrong_tokens)
dur_w = compute_uniform_durations(
    torch.tensor([wrong_tokens.shape[1]], device=tts.device),
    torch.tensor([T], device=tts.device))
wrong_expanded = length_regulate(wrong_emb, dur_w)[:, :T]

# Zero text (no conditioning)
zero_expanded = torch.zeros_like(correct_expanded)

# Run predictor with each
start = model.start_emb.expand(1, -1, -1)
input_emb = torch.cat([start, z[:, :-1]], dim=1)

with torch.no_grad():
    pred_correct = model.predictor(input_emb, correct_expanded)
    pred_wrong = model.predictor(input_emb, wrong_expanded)
    pred_zero = model.predictor(input_emb, zero_expanded)

    target = model.ema_proj_in(codec_emb.transpose(1, 2))

    # MSE with target
    mse_correct = torch.nn.functional.mse_loss(pred_correct, target).item()
    mse_wrong = torch.nn.functional.mse_loss(pred_wrong, target).item()
    mse_zero = torch.nn.functional.mse_loss(pred_zero, target).item()

    # How different are the predictions?
    diff_correct_wrong = torch.nn.functional.mse_loss(pred_correct, pred_wrong).item()
    diff_correct_zero = torch.nn.functional.mse_loss(pred_correct, pred_zero).item()

print(f'MSE vs target (correct text): {mse_correct:.6f}')
print(f'MSE vs target (wrong text):   {mse_wrong:.6f}')
print(f'MSE vs target (zero text):    {mse_zero:.6f}')
print(f'')
print(f'Diff between correct vs wrong pred: {diff_correct_wrong:.6f}')
print(f'Diff between correct vs zero pred:  {diff_correct_zero:.6f}')
print(f'')
print(f'Text influence: {(mse_wrong - mse_correct) / mse_correct * 100:.1f}% worse with wrong text')
" 2>&1

Test how much text conditioning affects predictions


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
v5 model loaded from output_v5_pred_dur/checkpoints/best.pt
MSE vs target (correct text): 0.273421
MSE vs target (wrong text):   0.272988
MSE vs target (zero text):    0.720276

Diff between correct vs wrong pred: 0.001564
Diff between correct vs zero pred:  0.480749

Text influence: -0.2% worse with wrong text


There it is. The numbers tell the whole story:

- **Correct text vs wrong text: 0.001564 MSE difference** — the model literally cannot tell the difference between correct Hindi and random English text
- **Zero text matters** (0.72 vs 0.27) — so it uses text as a "generic signal", but **doesn't read the actual content**
- Wrong text is -0.2% "worse" — basically identical to correct text

**The predictor learned to use text cross-attention as a generic bias signal but ignores actual text content.** It relies entirely on self-attention (previous audio frames) for prediction.

This is why AR synthesis just repeats the prompt — the text conditioning has no content-level influence.

**Root cause**: With teacher forcing + uniform duration alignment, the predictor can nearly perfectly predict the next frame from previous frames alone. Text becomes redundant during training.

**Fixes to consider**:
1. **Mask/drop audio frames during training** — force the model to rely on text (e.g., randomly zero out 30-50% of input audio frames)
2. **Scheduled sampling** — mix in AR-predicted frames during training so it can't rely on perfect previous frames
3. **Stronger text loss** — add an auxiliary loss that reconstructs text from predicted embeddings (forces text info into predictions)

Which direction do you want to try?


with current do we not doen these ?


The user is asking whether we already implemented any of these fixes in the current training. Let me check train_v5_pred.py more carefully - I already saw it has `--input_noise_max 0.05` which adds noise to input frames. Let me check if that's the same as masking/dropping frames.


    50→        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")
    51→
    52→    # Freeze proj_in and proj_out — they're already good
    53→    for p in model.proj_in.parameters():
    54→        p.requires_grad = False
    55→    for p in model.proj_out.parameters():
    56→        p.requires_grad = False
    57→    for p in model.ema_proj_in.parameters():
    58→        p.requires_grad = False
    59→
    60→    # Reinitialize predictor + text encoder + duration predictor from scratch
    61→    if args.reinit:
    62→        print("Reinitializing predictor + text encoder + duration predictor from scratch")
    63→        d_model = config.get("d_model", 256)
    64→        nhead = config.get("nhead", 4)
    65→        predictor_layers = config.get("predictor_layers", 6)
    66→        dropout = config.get("dropout", 0.1)
    67→        model.predictor = JEPAPredictor(
    68→            d_model=d_model, nhead=nhead,
    69→            num_layers=predictor_layers, dropout=dropout,
    70→        ).to(device)
    71→        model.text_encoder = TextEncoder(
    72→            vocab_size=config.get("text_vocab_size", 256),
    73→            d_model=d_model, nhead=nhead,
    74→            num_layers=config.get("text_encoder_layers", 4),
    75→            dropout=dropout,
    76→        ).to(device)
    77→        model.duration_predictor = DurationPredictor(
    78→            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
    79→        ).to(device)
    80→        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)
    81→
    82→    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    83→    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")
    84→
    85→    manifest_path = os.path.join(args.data_dir, "manifest.json")
    86→    loader, dataset = build_dataloader(
    87→        manifest_path, batch_size=args.batch_size,
    88→        num_workers=args.num_workers, max_codec_frames=args.max_codec_frames,
    89→    )
    90→    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")
    91→
    92→    # Only optimize trainable params
    93→    optimizer = torch.optim.AdamW(
    94→        [p for p in model.parameters() if p.requires_grad],
    95→        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    96→    )
    97→
    98→    total_steps = args.epochs * len(loader)
    99→    warmup_steps = min(2000, total_steps // 10)
   100→
   101→    def lr_lambda(step):
   102→        if step < warmup_steps:
   103→            return step / max(1, warmup_steps)
   104→        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
   105→        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())
   106→
   107→    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
   108→
   109→    ckpt_dir = Path(args.output_dir) / "checkpoints"
   110→    ckpt_dir.mkdir(parents=True, exist_ok=True)
   111→    writer = SummaryWriter(Path(args.output_dir) / "logs")
   112→
   113→    with open(Path(args.output_dir) / "config.json", "w") as f:
   114→        json.dump(config, f, indent=2)
   115→
   116→    model.train()
   117→    model.proj_in.eval()
   118→    model.proj_out.eval()
   119→    best_loss = float("inf")
   120→    global_step = 0
   121→
   122→    for epoch in range(args.epochs):
   123→        epoch_loss = 0.0
   124→        epoch_pred = 0.0
   125→        epoch_dur = 0.0
   126→        t0 = time.time()
   127→
   128→        for batch_idx, batch in enumerate(loader):
   129→            codec_emb = batch["codec_emb"].to(device)
   130→            text_tokens = batch["text_tokens"].to(device)
   131→            codec_mask = batch["codec_mask"].to(device)
   132→            text_mask = batch["text_mask"].to(device)
   133→
   134→            B = codec_emb.shape[0]
   135→            T = codec_emb.shape[2]
   136→            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]
   137→
   138→            # Text encoding (trainable)
   139→            text_emb = model.text_encoder(text_tokens, text_mask)
   140→
   141→            # Duration prediction
   142→            log_dur_pred = model.duration_predictor(text_emb, text_mask)
   143→
   144→            # Ground truth durations (uniform)
   145→            text_lengths = (~text_mask).sum(dim=1)
   146→            audio_lengths = (~codec_mask).sum(dim=1)
   147→            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
   148→            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))
   149→
   150→            # Duration loss
   151→            valid_text = (~text_mask).float()
   152→            dur_loss = F.mse_loss(
   153→                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
   154→            ) / (valid_text.sum() + 1e-8)
   155→
   156→            # Expand text to frame level using GT durations
   157→            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
   158→            if text_emb_expanded.shape[1] > T:
   159→                text_emb_expanded = text_emb_expanded[:, :T]
   160→            elif text_emb_expanded.shape[1] < T:
   161→                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
   162→                                device=device)
   163→                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)
   164→
   165→            # Project with FROZEN proj_in — targets are stable
   166→            with torch.no_grad():
   167→                z = model.proj_in(codec_seq)  # [B, T, d_model]
   168→                target_emb = z.clone()
   169→
   170→            # Prepend start, shift right
   171→            start = model.start_emb.expand(B, -1, -1)
   172→            input_emb = torch.cat([start, z[:, :-1]], dim=1)
   173→
   174→            # Input noise
   175→            noise_progress = min(1.0, global_step / (total_steps * 0.3))
   176→            noise_level = args.input_noise_max * noise_progress
   177→            if noise_level > 0:
   178→                noise = torch.zeros_like(input_emb)
   179→                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
   180→                                            device=device) * noise_level
   181→                input_emb = input_emb + noise
   182→
   183→            # Masks
   184→            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
   185→            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
   186→
   187→            # Predict with frame-level text
   188→            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)
   189→
   190→            # Prediction loss (MSE vs frozen proj_in targets)
   191→            valid = (~codec_mask).unsqueeze(-1)
   192→            pred_loss = F.mse_loss(
   193→                predicted * valid, target_emb * valid, reduction="sum"
   194→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   195→
   196→            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss
   197→
   198→            optimizer.zero_grad()
   199→            total_loss.backward()
   200→            grad_norm = torch.nn.utils.clip_grad_norm_(
   201→                [p for p in model.parameters() if p.requires_grad], args.grad_clip
   202→            )
   203→            optimizer.step()
   204→            scheduler.step()
   205→
   206→            epoch_loss += pred_loss.item()
   207→            epoch_pred += pred_loss.item()
   208→            epoch_dur += dur_loss.item()
   209→            global_step += 1
   210→
   211→            if global_step % args.log_every == 0:
   212→                writer.add_scalar("train/pred_loss", pred_loss.item(), global_step)
   213→                writer.add_scalar("train/dur_loss", dur_loss.item(), global_step)
   214→                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
   215→                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
   216→                writer.add_scalar("train/lr", optimizer.param_groups[0]["lr"], global_step)
   217→
   218→            if global_step % args.print_every == 0:
   219→                elapsed = time.time() - t0
   220→                print(
   221→                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
   222→                    f"pred={pred_loss.item():.4f} dur={dur_loss.item():.4f} | "
   223→                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"
   224→                )
   225→
   226→        avg_loss = epoch_loss / len(loader)
   227→        avg_dur = epoch_dur / len(loader)
   228→        epoch_time = time.time() - t0
   229→        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")
   230→
   231→        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
   232→            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
   233→            torch.save({
   234→                "epoch": epoch,
   235→                "global_step": global_step,
   236→                "model": model.state_dict(),
   237→                "optimizer": optimizer.state_dict(),
   238→                "scheduler": scheduler.state_dict(),
   239→                "config": config,
   240→                "loss": avg_loss,
   241→            }, ckpt_path)
   242→            print(f"  Saved: {ckpt_path}")
   243→
   244→        if avg_loss < best_loss:
   245→            best_loss = avg_loss
   246→            torch.save({
   247→                "epoch": epoch,
   248→                "global_step": global_step,
   249→                "model": model.state_dict(),
   250→                "config": config,
   251→                "loss": avg_loss,
   252→            }, ckpt_dir / "best.pt")
   253→            print(f"  New best (pred={avg_loss:.4f})")
   254→
   255→    writer.close()
   256→    print(f"\nDone. Best pred loss: {best_loss:.4f}")
   257→
   258→
   259→if __name__ == "__main__":
   260→    parser = argparse.ArgumentParser()
   261→    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
   262→    parser.add_argument("--data_dir", default="processed_data_codec")
   263→    parser.add_argument("--output_dir", default="output_v5_pred")
   264→    parser.add_argument("--pred_weight", type=float, default=1.0)
   265→    parser.add_argument("--dur_weight", type=float, default=1.0)
   266→    parser.add_argument("--input_noise_max", type=float, default=0.05)
   267→    parser.add_argument("--max_codec_frames", type=int, default=900)
   268→    parser.add_argument("--epochs", type=int, default=200)
   269→    parser.add_argument("--batch_size", type=int, default=256)
   270→    parser.add_argument("--lr", type=float, default=5e-5)
   271→    parser.add_argument("--grad_clip", type=float, default=1.0)
   272→    parser.add_argument("--num_workers", type=int, default=4)
   273→    parser.add_argument("--log_every", type=int, default=10)
   274→    parser.add_argument("--print_every", type=int, default=20)
   275→    parser.add_argument("--save_every", type=int, default=10)
   276→    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
   277→    args = parser.parse_args()
   278→    train(args)
   279→

        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

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

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

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

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

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

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

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

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

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

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

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

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)


We already have **input noise** (`input_noise_max=0.05`, ramped up over 30% of training) — but that's clearly not enough. It adds tiny Gaussian noise to audio frames, but the signal is still overwhelmingly there, so the predictor still just copies previous frames.

**What we're doing vs what's needed:**

| Technique | Current | What's needed |
|-----------|---------|---------------|
| Input noise | 0.05 (tiny) | Way too small to matter |
| Frame masking/dropout | **Not done** | Drop 30-50% of input audio frames to zeros — forces text reliance |
| Scheduled sampling | **Not done** | Replace GT frames with model's own predictions during training |
| Text-dependent loss | **Not done** | e.g., auxiliary CTC/alignment loss |

The noise at 0.05 is negligible — the embeddings have magnitude ~1.0, so it's a 5% perturbation. The model easily sees through that.

**Frame masking** is the simplest and most impactful fix — randomly zero out entire input audio frames so the predictor *must* read the text to know what to predict. Want me to add that and kick off a new run?


fucking do then we need correct thing .. fuck do


The user wants me to add frame masking/dropout to the training script and start a new training run. Let me do this quickly.

Key changes to train_v5_pred.py:
1. Add frame masking - randomly zero out 30-50% of input audio frames
2. This forces the predictor to rely on text conditioning

Let me also consider: should I kill the current run first? The user seems frustrated and wants the right thing done. Let me add the frame masking and start a new run.

I'll modify the training loop to mask input frames. I should add a `--frame_mask_ratio` argument.


Edit

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                # Don't mask start token (position 0)
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                if codec_mask is not None:
                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)


The file /home/ubuntu/lewm-tts/train_v5_pred.py has been updated successfully.

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                # Don't mask start token (position 0)
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                if codec_mask is not None:
                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)

"""
Train predictor + text encoder + duration predictor with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
Text is expanded to frame level via duration predictor — can't be ignored.
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import (
    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,
    DurationPredictor, compute_uniform_durations, length_regulate,
)
from dataset_v5 import build_dataloader

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

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["dur_weight"] = args.dur_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)

    # Load what we can from pretrained (projections, etc.)
    pretrained_state = ckpt["model"]
    model_state = model.state_dict()
    loaded_keys = []
    skipped_keys = []
    for k, v in pretrained_state.items():
        if k in model_state and model_state[k].shape == v.shape:
            model_state[k] = v
            loaded_keys.append(k)
        else:
            skipped_keys.append(k)
    model.load_state_dict(model_state)
    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")
    if skipped_keys:
        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

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

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

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

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

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

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

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

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

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

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

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

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)

+            # Frame masking — zero out random audio frames to force text reliance

+            if args.frame_mask_ratio > 0:

+                # Don't mask start token (position 0)

+                mask_probs = torch.rand(B, T, 1, device=device)

+                frame_mask = (mask_probs < args.frame_mask_ratio).float()

+                # Don't mask start embedding

+                # Don't mask padded positions

+                if codec_mask is not None:

+                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()

+                input_emb = input_emb * (1.0 - frame_mask)


    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--frame_mask_ratio", type=float, default=0.0,
                        help="Fraction of input audio frames to zero out (forces text reliance)")


    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--frame_mask_ratio", type=float, default=0.0,
                        help="Fraction of input audio frames to zero out (forces text reliance)")

"""
Train predictor + text encoder + duration predictor with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
Text is expanded to frame level via duration predictor — can't be ignored.
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import (
    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,
    DurationPredictor, compute_uniform_durations, length_regulate,
)
from dataset_v5 import build_dataloader

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

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["dur_weight"] = args.dur_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)

    # Load what we can from pretrained (projections, etc.)
    pretrained_state = ckpt["model"]
    model_state = model.state_dict()
    loaded_keys = []
    skipped_keys = []
    for k, v in pretrained_state.items():
        if k in model_state and model_state[k].shape == v.shape:
            model_state[k] = v
            loaded_keys.append(k)
        else:
            skipped_keys.append(k)
    model.load_state_dict(model_state)
    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")
    if skipped_keys:
        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

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

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

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

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

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

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

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

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

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

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

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

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                # Don't mask start token (position 0)
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                if codec_mask is not None:
                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)

+    parser.add_argument("--frame_mask_ratio", type=float, default=0.0,

+                        help="Fraction of input audio frames to zero out (forces text reliance)")


Now kill the current run and start a new one with frame masking:


Bash

kill 3768655 2>/dev/null; sleep 1; ps aux | grep train_v5_pred | grep -v grep | awk '{print $2}' | xargs kill 2>/dev/null; sleep 1; ps aux | grep train_v5_pred | grep -v grep

Kill current pred_dur training processes


(Bash completed with no output)


Good, all killed. Now start fresh with frame masking at 0.4 (40% of frames masked):


mkdir -p output_v5_pred_mask && rm -rf output_v5_pred_mask/logs output_v5_pred_mask/checkpoints 2>/dev/null && nohup python3 train_v5_pred.py \
  --pretrained output_v5/checkpoints/epoch_0120.pt \
  --output_dir output_v5_pred_mask \
  --reinit \
  --pred_weight 1.0 \
  --dur_weight 1.0 \
  --lr 5e-5 \
  --batch_size 128 \
  --epochs 200 \
  --print_every 20 \
  --save_every 10 \
  --input_noise_max 0.05 \
  --frame_mask_ratio 0.4 > output_v5_pred_mask/train.log 2>&1 & echo "PID: $!"

Start training with 40% frame masking


PID: 3783232


sleep 15 && tail -20 output_v5_pred_mask/train.log

Wait and check training started


Device: cuda
Loaded 144 params from output_v5/checkpoints/epoch_0120.pt (skipped 36)
  Skipped: ['predictor.transformer.layers.0.multihead_attn.in_proj_weight', 'predictor.transformer.layers.0.multihead_attn.in_proj_bias', 'predictor.transformer.layers.0.multihead_attn.out_proj.weight', 'predictor.transformer.layers.0.multihead_attn.out_proj.bias', 'predictor.transformer.layers.0.norm3.weight', 'predictor.transformer.layers.0.norm3.bias', 'predictor.transformer.layers.1.multihead_attn.in_proj_weight', 'predictor.transformer.layers.1.multihead_attn.in_proj_bias', 'predictor.transformer.layers.1.multihead_attn.out_proj.weight', 'predictor.transformer.layers.1.multihead_attn.out_proj.bias']...
Reinitializing predictor + text encoder + duration predictor from scratch
Trainable params: 8.69M (predictor + text encoder + dur predictor + start_emb)
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 53 batches/epoch
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/train_v5_pred.py", line 292, in <module>
    train(args)
  File "/home/ubuntu/lewm-tts/train_v5_pred.py", line 192, in train
    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
UnboundLocalError: local variable 'pred_mask' referenced before assignment


The mask variable is defined a few lines later. Let me fix the ordering:


Edit

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

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                # Don't mask start token (position 0)
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                if codec_mask is not None:
                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

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

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding (position 0)
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)


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

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding (position 0)
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)

"""
Train predictor + text encoder + duration predictor with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
Text is expanded to frame level via duration predictor — can't be ignored.
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import (
    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,
    DurationPredictor, compute_uniform_durations, length_regulate,
)
from dataset_v5 import build_dataloader

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

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["dur_weight"] = args.dur_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)

    # Load what we can from pretrained (projections, etc.)
    pretrained_state = ckpt["model"]
    model_state = model.state_dict()
    loaded_keys = []
    skipped_keys = []
    for k, v in pretrained_state.items():
        if k in model_state and model_state[k].shape == v.shape:
            model_state[k] = v
            loaded_keys.append(k)
        else:
            skipped_keys.append(k)
    model.load_state_dict(model_state)
    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")
    if skipped_keys:
        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

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

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

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

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

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

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

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

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

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

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

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

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

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

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Frame masking — zero out random audio frames to force text reliance
            if args.frame_mask_ratio > 0:
                # Don't mask start token (position 0)
                mask_probs = torch.rand(B, T, 1, device=device)
                frame_mask = (mask_probs < args.frame_mask_ratio).float()
                # Don't mask start embedding
                frame_mask[:, 0] = 0.0
                # Don't mask padded positions
                if codec_mask is not None:
                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()
                input_emb = input_emb * (1.0 - frame_mask)

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

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

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

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--frame_mask_ratio", type=float, default=0.0,
                        help="Fraction of input audio frames to zero out (forces text reliance)")
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)

+            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)

+            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

-                # Don't mask start token (position 0)

-                # Don't mask start embedding

+                # Don't mask start embedding (position 0)

-                if codec_mask is not None:

-                    frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()

+                frame_mask = frame_mask * (~pred_mask).unsqueeze(-1).float()

-            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)

-            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)


Bash

rm -rf output_v5_pred_mask/logs output_v5_pred_mask/checkpoints 2>/dev/null && nohup python3 train_v5_pred.py \
  --pretrained output_v5/checkpoints/epoch_0120.pt \
  --output_dir output_v5_pred_mask \
  --reinit \
  --pred_weight 1.0 \
  --dur_weight 1.0 \
  --lr 5e-5 \
  --batch_size 128 \
  --epochs 200 \
  --print_every 20 \
  --save_every 10 \
  --input_noise_max 0.05 \
  --frame_mask_ratio 0.4 > output_v5_pred_mask/train.log 2>&1 & echo "PID: $!"

Restart training with frame masking fix


PID: 3783616


sleep 40 && tail -15 output_v5_pred_mask/train.log

Check training is running


Device: cuda
Loaded 144 params from output_v5/checkpoints/epoch_0120.pt (skipped 36)
  Skipped: ['predictor.transformer.layers.0.multihead_attn.in_proj_weight', 'predictor.transformer.layers.0.multihead_attn.in_proj_bias', 'predictor.transformer.layers.0.multihead_attn.out_proj.weight', 'predictor.transformer.layers.0.multihead_attn.out_proj.bias', 'predictor.transformer.layers.0.norm3.weight', 'predictor.transformer.layers.0.norm3.bias', 'predictor.transformer.layers.1.multihead_attn.in_proj_weight', 'predictor.transformer.layers.1.multihead_attn.in_proj_bias', 'predictor.transformer.layers.1.multihead_attn.out_proj.weight', 'predictor.transformer.layers.1.multihead_attn.out_proj.bias']...
Reinitializing predictor + text encoder + duration predictor from scratch
Trainable params: 8.69M (predictor + text encoder + dur predictor + start_emb)
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 53 batches/epoch
  [1/200] step 20 | pred=11.5016 dur=1.9359 | grad=60.14 | 13.6s
  [1/200] step 40 | pred=11.3191 dur=1.0257 | grad=40.23 | 26.3s
Epoch 1/200 | pred=11.3326 dur=1.5860 | 34.4s
  New best (pred=11.3326)


Running. pred_loss starts at 11.33 (much higher than the old run's ~0.6 start) — that's expected because with 40% frames masked, the model can't just copy previous frames anymore. It's being **forced to learn from text**.

**New run**: `output_v5_pred_mask/` — PID 3783616
- 40% frame masking + input noise 0.05
- ~34s/epoch, 200 epochs
- pred_loss should converge slower but with **real text conditioning**

The high initial loss is a good sign — it means the model actually needs to work harder and can't cheat by copying audio frames.


stattus


Check training progress


  New best (pred=5.7257)
  [9/200] step 440 | pred=4.6710 dur=0.0548 | grad=6.13 | 10.5s
  [9/200] step 460 | pred=4.2612 dur=0.0641 | grad=6.01 | 23.0s
Epoch 9/200 | pred=4.4092 dur=0.0571 | 33.9s
  New best (pred=4.4092)
  [10/200] step 480 | pred=3.7705 dur=0.0572 | grad=5.68 | 2.2s
  [10/200] step 500 | pred=3.2859 dur=0.0514 | grad=5.29 | 14.9s
  [10/200] step 520 | pred=2.8070 dur=0.0523 | grad=4.88 | 27.8s
Epoch 10/200 | pred=3.1615 dur=0.0539 | 34.2s
  Saved: output_v5_pred_mask/checkpoints/epoch_0010.pt
  New best (pred=3.1615)
  [11/200] step 540 | pred=2.4054 dur=0.0412 | grad=4.33 | 6.6s
  [11/200] step 560 | pred=2.0894 dur=0.0520 | grad=4.10 | 19.1s
  [11/200] step 580 | pred=1.8066 dur=0.0444 | grad=3.48 | 31.9s
Epoch 11/200 | pred=2.1493 dur=0.0516 | 33.8s
  New best (pred=2.1493)
  [12/200] step 600 | pred=1.5365 dur=0.0505 | grad=3.05 | 11.0s
  [12/200] step 620 | pred=1.3393 dur=0.0435 | grad=2.34 | 23.7s
Epoch 12/200 | pred=1.4551 dur=0.0494 | 34.0s
  New best (pred=1.4551)


**Epoch 12/200** — pred loss dropping fast: 11.3 → 1.45 in 12 epochs. Still hitting new bests every epoch. Duration loss already converged at 0.05.

The loss is still much higher than the old run (which was at 0.42 by ep54) — but that's the point. With 40% frames masked, the model can't cheat. It has to actually learn text→audio mapping to fill the gaps.


run samples


Check current epoch