Write

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        # Speaker conditioning
        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)  # [1, d]
            # FiLM on text embeddings (matches training)
            text_emb = self.model._apply_speaker_film(
                text_emb, spk_cond, self.model.latent_film
            )

        # Use learned start embedding
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

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

        # FiLM on predicted latent (matches training where z gets FiLM'd)
        emb_for_decode = self.model._apply_speaker_film(
            next_emb, spk_cond, self.model.latent_film
        ) if spk_cond is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed FiLM-conditioned embedding as input (matches training input_emb)
            next_emb, cache = self.model.predict_next_cached(
                emb_for_decode, text_emb, step, cache
            )
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            emb_for_decode = self.model._apply_speaker_film(
                next_emb, spk_cond, self.model.latent_film
            ) if spk_cond is not None else next_emb
            all_embeddings.append(emb_for_decode)

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

        audio_embs = torch.cat(all_embeddings, dim=1)
        # Pass speaker conditioning to mel decoder (FiLM-conditioned reconstruction)
        mel = self.model.mel_decoder(audio_embs, spk_cond)
        waveform = self.vocos_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

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

        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)

        z, mu, logvar = self.model.audio_encoder(mel)
        # FiLM condition the latent before decoding
        if spk_cond is not None:
            mu = self.model._apply_speaker_film(mu, spk_cond, self.model.latent_film)
        mel_recon = self.model.mel_decoder(mu, spk_cond)

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

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

        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)

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

        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)
        if spk_cond is not None:
            text_emb = self.model._apply_speaker_film(
                text_emb, spk_cond, self.model.latent_film
            )

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

        if spk_cond is not None:
            combined = self.model._apply_speaker_film(
                combined, spk_cond, self.model.latent_film
            )
        mel_recon = self.model.mel_decoder(combined, spk_cond)
        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


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

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        # Speaker conditioning
        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)  # [1, d]
            # FiLM on text embeddings (matches training)
            text_emb = self.model._apply_speaker_film(
                text_emb, spk_cond, self.model.latent_film
            )

        # Use learned start embedding
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

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

        # FiLM on predicted latent (matches training where z gets FiLM'd)
        emb_for_decode = self.model._apply_speaker_film(
            next_emb, spk_cond, self.model.latent_film
        ) if spk_cond is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed FiLM-conditioned embedding as input (matches training input_emb)
            next_emb, cache = self.model.predict_next_cached(
                emb_for_decode, text_emb, step, cache
            )
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            emb_for_decode = self.model._apply_speaker_film(
                next_emb, spk_cond, self.model.latent_film
            ) if spk_cond is not None else next_emb
            all_embeddings.append(emb_for_decode)

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

        audio_embs = torch.cat(all_embeddings, dim=1)
        # Pass speaker conditioning to mel decoder (FiLM-conditioned reconstruction)
        mel = self.model.mel_decoder(audio_embs, spk_cond)
        waveform = self.vocos_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

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

        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)

        z, mu, logvar = self.model.audio_encoder(mel)
        # FiLM condition the latent before decoding
        if spk_cond is not None:
            mu = self.model._apply_speaker_film(mu, spk_cond, self.model.latent_film)
        mel_recon = self.model.mel_decoder(mu, spk_cond)

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

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

        spk_cond = None
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            spk_cond = self.model.speaker_embed(spk_id)

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

        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)
        if spk_cond is not None:
            text_emb = self.model._apply_speaker_film(
                text_emb, spk_cond, self.model.latent_film
            )

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

        if spk_cond is not None:
            combined = self.model._apply_speaker_film(
                combined, spk_cond, self.model.latent_film
            )
        mel_recon = self.model.mel_decoder(combined, spk_cond)
        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

-        # Add speaker conditioning to text embeddings

-            text_emb = self.model._add_speaker(text_emb, spk_id)

-            # Cache speaker embedding for adding to predicted latents

-            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]

+            spk_cond = self.model.speaker_embed(spk_id)  # [1, d]

+            # FiLM on text embeddings (matches training)

+            text_emb = self.model._apply_speaker_film(

+                text_emb, spk_cond, self.model.latent_film

-        # Use learned start embedding instead of random noise

-        # Add speaker embedding to predicted latent (matches training where z = z + spk)

-        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb

+        # FiLM on predicted latent (matches training where z gets FiLM'd)

+        emb_for_decode = self.model._apply_speaker_film(

+            next_emb, spk_cond, self.model.latent_film

+        ) if spk_cond is not None else next_emb

-            # Feed speaker-conditioned embedding as input (matches training input_emb)

-            next_input = emb_for_decode

-            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)

+            # Feed FiLM-conditioned embedding as input (matches training input_emb)

+            next_emb, cache = self.model.predict_next_cached(

+                emb_for_decode, text_emb, step, cache

-        mel = self.model.mel_decoder(audio_embs)

+        # Pass speaker conditioning to mel decoder (FiLM-conditioned reconstruction)

+        mel = self.model.mel_decoder(audio_embs, spk_cond)

+    def reconstruct(self, mel_path, speaker_id=None):

+        if self.model.speaker_embed is not None and speaker_id is not None:

+            spk_id = torch.tensor([speaker_id], device=self.device)

-        mel_recon = self.model.mel_decoder(mu)

+        # FiLM condition the latent before decoding

+            mu = self.model._apply_speaker_film(mu, spk_cond, self.model.latent_film)

+        mel_recon = self.model.mel_decoder(mu, spk_cond)

-    def reconstruct_with_prediction(self, mel_path):

+    def reconstruct_with_prediction(self, mel_path, speaker_id=None):

-        mel_recon = self.model.mel_decoder(combined)

+            combined = self.model._apply_speaker_film(

+                combined, spk_cond, self.model.latent_film

+        mel_recon = self.model.mel_decoder(combined, spk_cond)

-        waveform, sr = tts.reconstruct(args.mel)

+        waveform, sr = tts.reconstruct(args.mel, speaker_id=args.speaker_id)

-        waveform, sr = tts.reconstruct_with_prediction(args.mel)

+        waveform, sr = tts.reconstruct_with_prediction(args.mel, speaker_id=args.speaker_id)

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

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

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)

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

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

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

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

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

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

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

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

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

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

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


Now update train.py to log the new cosine loss:


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

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.to(device)

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Update EMA target encoder
            model.update_ema()

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

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

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

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

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

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

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

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

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


Add `cosine_weight` to the config and log the new losses:


Edit

        "input_noise": 0.0,  # will be annealed during training
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,

        "input_noise": 0.0,  # will be annealed during training
        "cosine_weight": args.cosine_weight,
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,


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

        "input_noise": 0.0,  # will be annealed during training
        "cosine_weight": args.cosine_weight,
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.to(device)

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Update EMA target encoder
            model.update_ema()

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

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

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

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

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

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

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

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

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

+        "cosine_weight": args.cosine_weight,


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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            mse_loss = losses["mse_loss"]
            cosine_loss = losses["cosine_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_loss"]
            speaker_loss = losses["speaker_loss"]


            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            mse_loss = losses["mse_loss"]
            cosine_loss = losses["cosine_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_loss"]
            speaker_loss = losses["speaker_loss"]

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "cosine_weight": args.cosine_weight,
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.to(device)

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Update EMA target encoder
            model.update_ema()

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

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

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

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

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

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

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

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

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

+            mse_loss = losses["mse_loss"]

+            cosine_loss = losses["cosine_loss"]


                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
                writer.add_scalar("train/prediction_loss", pred_loss.item(), global_step)
                writer.add_scalar("train/kl_loss", kl_loss.item(), global_step)

                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
                writer.add_scalar("train/prediction_loss", pred_loss.item(), global_step)
                writer.add_scalar("train/mse_loss", mse_loss.item(), global_step)
                writer.add_scalar("train/cosine_loss", cosine_loss.item(), global_step)
                writer.add_scalar("train/kl_loss", kl_loss.item(), global_step)


                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
                writer.add_scalar("train/prediction_loss", pred_loss.item(), global_step)
                writer.add_scalar("train/mse_loss", mse_loss.item(), global_step)
                writer.add_scalar("train/cosine_loss", cosine_loss.item(), global_step)
                writer.add_scalar("train/kl_loss", kl_loss.item(), global_step)

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "cosine_weight": args.cosine_weight,
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.to(device)

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            mse_loss = losses["mse_loss"]
            cosine_loss = losses["cosine_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_loss"]
            speaker_loss = losses["speaker_loss"]

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

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

            optimizer.step()
            scheduler.step()

            # Update EMA target encoder
            model.update_ema()

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

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

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

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

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

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

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

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

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

+                writer.add_scalar("train/mse_loss", mse_loss.item(), global_step)

+                writer.add_scalar("train/cosine_loss", cosine_loss.item(), global_step)


    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)

    parser.add_argument("--cosine_weight", type=float, default=1.0)
    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)


    parser.add_argument("--cosine_weight", type=float, default=1.0)
    parser.add_argument("--n_speakers", type=int, default=1)
    parser.add_argument("--speaker_weight", type=float, default=1.0)

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
        "pred_weight": args.pred_weight,
        "ema_decay": args.ema_decay,
        "free_bits": args.free_bits,
        "input_noise": 0.0,  # will be annealed during training
        "cosine_weight": args.cosine_weight,
        "n_speakers": args.n_speakers,
        "speaker_weight": args.speaker_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)
            speaker_id = batch.get("speaker_id")
            if speaker_id is not None:
                speaker_id = speaker_id.to(device)

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

            # Input noise annealing: ramp up over first 30% of training
            # Bridges train/inference gap by simulating imperfect AR inputs
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

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

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            mse_loss = losses["mse_loss"]
            cosine_loss = losses["cosine_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]
            spectral_loss = losses["spectral_loss"]
            speaker_loss = losses["speaker_loss"]

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

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

            optimizer.step()
            scheduler.step()

            # Update EMA target encoder
            model.update_ema()

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

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

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

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

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

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

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

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

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

+    parser.add_argument("--cosine_weight", type=float, default=1.0)


Now run the comprehensive test:


Bash

python3 -c "
import torch
from model import build_model, count_parameters

print('='*60)
print('TEST 1: Single-speaker backward compatibility')
print('='*60)
model_s, cfg_s = build_model()
B = 2
mel = torch.randn(B, 100, 400)
text = torch.randint(0, 256, (B, 20))
mel_mask = torch.zeros(B, 400, dtype=torch.bool)
text_mask = torch.zeros(B, 20, dtype=torch.bool)

model_s.train()
losses = model_s(mel, text, mel_mask, text_mask)
losses['total_loss'].backward()
print('Forward+backward OK')
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')
model_s.zero_grad()

print()
print('='*60)
print('TEST 2: Multi-speaker (10 speakers)')
print('='*60)
config_ms = {
    'd_model': 256, 'nhead': 4, 'n_mels': 100,
    'text_vocab_size': 256, 'text_encoder_layers': 4,
    'audio_encoder_layers': 4, 'predictor_layers': 6,
    'dropout': 0.1, 'kl_weight': 0.05, 'pred_weight': 10.0,
    'cosine_weight': 1.0,
    'ema_decay': 0.998, 'free_bits': 2.0, 'input_noise': 0.1,
    'n_speakers': 10, 'speaker_weight': 1.0,
    'recon_weight': 1.0, 'spectral_weight': 0.5,
}
model_ms, _ = build_model(config_ms)
model_ms.train()

B = 4
mel = torch.randn(B, 100, 400)
text = torch.randint(0, 256, (B, 30))
mel_mask = torch.zeros(B, 400, dtype=torch.bool)
text_mask = torch.zeros(B, 30, dtype=torch.bool)
spk = torch.tensor([0, 3, 7, 9])

losses = model_ms(mel, text, mel_mask, text_mask, speaker_id=spk)
losses['total_loss'].backward()
print('Forward+backward OK')
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')

# Check gradient flow
print()
print('Gradient flow check:')
print(f'  AudioEncoder proj_mu:     {model_ms.audio_encoder.proj_mu.weight.grad.norm():.4f}')
print(f'  Speaker embed:            {model_ms.speaker_embed.weight.grad.norm():.4f}')
print(f'  Speaker classifier[0]:    {model_ms.speaker_classifier[0].weight.grad.norm():.4f}')
print(f'  Latent FiLM proj:         {model_ms.latent_film.proj.weight.grad.norm():.4f}')
print(f'  Target FiLM proj:         {model_ms.target_film.proj.weight.grad.norm():.4f}')
print(f'  MelDecoder input FiLM:    {model_ms.mel_decoder.input_film.proj.weight.grad.norm():.4f}')
print(f'  MelDecoder res1 FiLM:     {model_ms.mel_decoder.res1.film.proj.weight.grad.norm():.4f}')
print(f'  EMA target proj:          {model_ms.ema_target_proj.weight.grad.norm():.4f}')

# Non-zero rows in speaker embed (should be 4 = batch speakers)
grad_spk = model_ms.speaker_embed.weight.grad
nonzero = (grad_spk.abs().sum(dim=1) > 0).sum().item()
print(f'  Speaker embed active rows: {nonzero}/10 (expect 4)')
model_ms.zero_grad()

print()
print('='*60)
print('TEST 3: FiLM actually modulates (not identity)')
print('='*60)
model_ms.eval()
with torch.no_grad():
    mel_same = mel[:1].expand(2, -1, -1)
    text_same = text[:1].expand(2, -1)
    spk_diff = torch.tensor([0, 7])

    spk_cond = model_ms.speaker_embed(spk_diff)

    # Raw encoder output (same for both)
    z, mu, _ = model_ms.audio_encoder(mel_same)

    # FiLM with different speakers
    z_0 = model_ms.latent_film(z[0:1], spk_cond[0:1])
    z_1 = model_ms.latent_film(z[1:2], spk_cond[1:2])

    raw_diff = (z[0] - z[1]).abs().mean().item()
    film_diff = (z_0[0] - z_1[0]).abs().mean().item()
    print(f'  Raw encoder diff (same mel): {raw_diff:.6f} (should be ~0)')
    print(f'  FiLM diff (diff speakers):   {film_diff:.4f} (should be >0)')

    # FiLM gamma/beta check
    params = model_ms.latent_film.proj(spk_cond[0])
    gamma, beta = params.chunk(2)
    print(f'  FiLM gamma mean: {gamma.mean():.4f} (init ~1.0)')
    print(f'  FiLM beta mean:  {beta.mean():.4f} (init ~0.0)')

print()
print('='*60)
print('TEST 4: Multi-scale targets')
print('='*60)
model_ms.eval()
with torch.no_grad():
    target = model_ms._compute_multiscale_targets(mel[:1], None)
    print(f'  Target shape: {target.shape} (should be [1, T_down, 256])')
    print(f'  Target norm: {target.norm():.4f} (should be reasonable)')

print()
print('='*60)
print('TEST 5: Variable-length padding')
print('='*60)
model_ms.train()
mel_mask2 = torch.zeros(B, 400, dtype=torch.bool)
mel_mask2[0, 300:] = True
mel_mask2[1, 200:] = True
text_mask2 = torch.zeros(B, 30, dtype=torch.bool)
text_mask2[2, 20:] = True
losses = model_ms(mel, text, mel_mask2, text_mask2, speaker_id=spk)
losses['total_loss'].backward()
print('Padded forward+backward OK')
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')

print()
print('='*60)
print('MODEL SIZE SUMMARY')
print('='*60)
total = sum(p.numel() for p in model_ms.parameters())
trainable = count_parameters(model_ms)
ema = sum(p.numel() for p in model_ms.ema_audio_encoder.parameters())
pred = sum(p.numel() for p in model_ms.predictor.parameters())
cls = sum(p.numel() for p in model_ms.speaker_classifier.parameters())

print(f'Total params:     {total/1e6:.2f}M')
print(f'Trainable:        {trainable/1e6:.2f}M')
print(f'EMA (frozen):     {ema/1e6:.2f}M')
print(f'Predictor:        {pred/1e6:.2f}M')
print()

# Inference params = total - EMA - predictor - classifier
inf_params = total - ema - pred - cls
print(f'Inference params: {inf_params/1e6:.2f}M')
print(f'FP16 disk size:   ~{total*2/1e6:.1f}MB (full) / ~{inf_params*2/1e6:.1f}MB (inference)')

print()
print('Component breakdown:')
for name, mod in [
    ('TextEncoder', model_ms.text_encoder),
    ('AudioEncoder', model_ms.audio_encoder),
    ('Predictor', model_ms.predictor),
    ('MelDecoder', model_ms.mel_decoder),
    ('EMA AudioEncoder', model_ms.ema_audio_encoder),
    ('EMA target proj', model_ms.ema_target_proj),
    ('Speaker Embed', model_ms.speaker_embed),
    ('Latent FiLM', model_ms.latent_film),
    ('Target FiLM', model_ms.target_film),
    ('Speaker Classifier', model_ms.speaker_classifier),
]:
    n = sum(p.numel() for p in mod.parameters())
    print(f'  {name:25s}: {n/1e6:.3f}M ({n:,} params)')

print()
print('ALL TESTS PASSED')
" 2>&1

Comprehensive test of full rewrite