        self.config = config

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

        # Compute start token from training data statistics

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

        self.config = config

        # Compute start token from training data statistics
        # Use a random normal scaled to match training embedding norms
        self._start_emb_scale = 5.0  # approximate from diagnostics

        print(f"Model loaded from {checkpoint_path}")

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        waveform = self.vocoder.decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """
        Teacher-forced reconstruction: encode real mel → decode back.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        waveform = self.vocoder.decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """
        Hybrid: encode real mel, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        mel_recon = self.model.mel_decoder(combined)
        waveform = self.vocoder.decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    def save_audio(self, waveform, sr, output_path):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

+        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""

+        import torch.nn.functional as F

+        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)

+        return self._vocos.decode(mel_padded)


Now replace all `.vocoder.decode(` calls with `.vocos_decode(`:


The file /home/ubuntu/lewm-tts/inference.py has been updated. All occurrences were successfully replaced.

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

        self.config = config

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

        # Compute start token from training data statistics
        # Use a random normal scaled to match training embedding norms
        self._start_emb_scale = 5.0  # approximate from diagnostics

        print(f"Model loaded from {checkpoint_path}")

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        waveform = self.vocoder.decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """
        Teacher-forced reconstruction: encode real mel → decode back.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        waveform = self.vocoder.decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """
        Hybrid: encode real mel, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        mel_recon = self.model.mel_decoder(combined)
        waveform = self.vocoder.decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    def save_audio(self, waveform, sr, output_path):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

-        waveform = self.vocoder.decode(mel.cpu())

+        waveform = self.vocos_decode(mel.cpu())

-        waveform = self.vocoder.decode(mel_recon.cpu())

+        waveform = self.vocos_decode(mel_recon.cpu())


Now run all 3 modes:


Bash

python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --mel processed_data/mels/000000.pt --mode reconstruct --output output/v1_e350_vocos_reconstruct.wav

v1 reconstruct with Vocos


python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --mel processed_data/mels/000000.pt --mode predict --output output/v1_e350_vocos_predict.wav

v1 predict with Vocos


python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v1_e350_vocos_synth.wav --max_steps 500

v1 synthesize Hindi with Vocos


Saved: output/v1_e350_vocos_reconstruct.wav (9.50s)


Saved: output/v1_e350_vocos_predict.wav (9.50s)


Error: Exit code 1
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/inference.py", line 205, in <module>
    main()
  File "/home/ubuntu/lewm-tts/inference.py", line 193, in main
    waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/lewm-tts/inference.py", line 93, in synthesize
    start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
AttributeError: 'LeWMTTSInference' object has no attribute '_start_emb_scale'


The vocos_decode method broke the comment continuation. Let me fix:


    64→
    65→    def vocos_decode(self, mel):
    66→        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
    67→        import torch.nn.functional as F
    68→        # Pad 80 → 100 mels
    69→        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
    70→        return self._vocos.decode(mel_padded)
    71→
    72→        # Compute start token from training data statistics
    73→        # Use a random normal scaled to match training embedding norms
    74→        self._start_emb_scale = 5.0  # approximate from diagnostics
    75→
    76→        print(f"Model loaded from {checkpoint_path}")
    77→
    78→    def text_to_tokens(self, text):

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

        # Compute start token from training data statistics
        # Use a random normal scaled to match training embedding norms
        self._start_emb_scale = 5.0  # approximate from diagnostics

        print(f"Model loaded from {checkpoint_path}")

    def text_to_tokens(self, text):


The `_start_emb_scale` and print got orphaned after the return. Let me fix:


Edit

        return self._vocos.decode(mel_padded)

        # Compute start token from training data statistics
        # Use a random normal scaled to match training embedding norms
        self._start_emb_scale = 5.0  # approximate from diagnostics

        print(f"Model loaded from {checkpoint_path}")


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

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

        self.config = config

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

        # Compute start token from training data statistics
        # Use a random normal scaled to match training embedding norms
        self._start_emb_scale = 5.0  # approximate from diagnostics

        print(f"Model loaded from {checkpoint_path}")

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        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.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        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, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        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):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

-        # Compute start token from training data statistics

-        # Use a random normal scaled to match training embedding norms

-        self._start_emb_scale = 5.0  # approximate from diagnostics

-        print(f"Model loaded from {checkpoint_path}")


Now add the missing init lines back into `__init__`:


        self.config = config

    def vocos_decode

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

    def vocos_decode


        self.config = config

    def vocos_decode

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

    def vocos_decode

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

        self.config = config

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        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.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        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, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        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):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

+        print(f"Model loaded from {checkpoint_path}")


Bash

python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v1_e350_vocos_synth.wav --max_steps 500

v1 synthesize Hindi with Vocos


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_e350_vocos_synth.wav (2.34s)


Now v2:


cd /home/ubuntu/lewm-tts/v2 && python3 inference.py --checkpoint output/checkpoints/epoch_0150.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v2_e150_synth.wav --max_steps 500

v2 synthesize Hindi with DAC vocoder


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loading DAC 24kHz...
Model loaded. Start scale: 135.694
Saved: output/v2_e150_synth.wav (0.55s)


All done. Here's the final output with Vocos for v1:

**v1 (mel + Vocos) — epoch 350:**
| Mode | File | Duration |
|------|------|----------|
| Reconstruct | `output/v1_e350_vocos_reconstruct.wav` | 9.50s |
| Predict (teacher-forced) | `output/v1_e350_vocos_predict.wav` | 9.50s |
| Synthesize (Hindi) | `output/v1_e350_vocos_synth.wav` | 2.34s |

**v2 (DAC) — epoch 150:**
| Mode | File | Duration |
|------|------|----------|
| Synthesize (Hindi) | `v2/output/v2_e150_synth.wav` | 0.55s |

Fixed the vocoder — now using **Vocos** instead of Griffin-Lim for v1. The reconstruct/predict modes produce good length audio. The pure AR synthesis is still short on both versions, which suggests the AR generation is stopping too early or hasn't learned long-form generation yet.


 Predict (teacher-forced) │ output/v1_e350_vocos_predict.wav     │ 9.50s    │
  ├──────────────────────────┼──────────────────────────────────────┼──────────┤
  │ Synthesize (Hindi)  wt the differnce between both


**Reconstruct** = encode real audio mel → decode it back. Tests if the mel decoder + vocoder can reproduce audio faithfully. No prediction involved — just autoencoder roundtrip.

**Predict (teacher-forced)** = encode real audio mel → run the JEPA predictor on the real embeddings (it sees all real frames, predicts the next one) → decode predicted embeddings. Tests if the predictor can accurately reproduce embeddings when given ground truth context.

**Synthesize** = pure AR generation from text only. No real audio input at all — the model generates each frame from its own previous predictions. This is the actual TTS mode, but also hardest since errors compound.

So quality ranking is typically: Reconstruct > Predict > Synthesize. The first two use real audio as a crutch — only Synthesize is true text-to-speech.


Synthesize - this is noise ..at 25th/75th epoch we goyt good result but not here... u dng infernce something worng and for v2 also


The user says at epoch 25/75 we got good synthesis results but now at epoch 350 it's noise. This could be overfitting, or the inference code changed. Let me check if there's an older inference script or if something changed. Let me also check the training log to see if losses diverged.

Let me also check git history for changes to inference.py.


Let me check what changed — could be overfitting at later epochs or the inference code got modified.


cd /home/ubuntu/lewm-tts && git log --oneline -20 2>/dev/null || echo "no git"

Check git history


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
# Check losses at different epochs
for ep in ['epoch_0070', 'epoch_0100', 'epoch_0350', 'best']:
    ckpt = torch.load(f'output/checkpoints/{ep}.pt', map_location='cpu', weights_only=False)
    print(f'{ep}: epoch={ckpt.get(\"epoch\",\"?\")}, losses={ckpt.get(\"losses\",\"?\")}, best_loss={ckpt.get(\"best_loss\",\"?\")}')
" 2>&1

Check losses across checkpoints


tail -100 /home/ubuntu/lewm-tts/output/train_v1_resumed.log 2>/dev/null | head -50

Check recent training log


no git


epoch_0070: epoch=69, losses=?, best_loss=?
epoch_0100: epoch=99, losses=?, best_loss=?
epoch_0350: epoch=349, losses=?, best_loss=?
best: epoch=368, losses=?, best_loss=?


  [358/1000] step 38560 | loss=0.4660 pred=0.0467 kl=1.2416 recon=0.2951 | grad=1.36 | 1.3s
  [358/1000] step 38580 | loss=0.4650 pred=0.0466 kl=1.2427 recon=0.2941 | grad=1.15 | 10.7s
  [358/1000] step 38600 | loss=0.4656 pred=0.0467 kl=1.2423 recon=0.2947 | grad=0.61 | 19.9s
  [358/1000] step 38620 | loss=0.4654 pred=0.0467 kl=1.2437 recon=0.2943 | grad=0.99 | 29.3s
  [358/1000] step 38640 | loss=0.4647 pred=0.0465 kl=1.2428 recon=0.2939 | grad=0.70 | 38.4s
  [358/1000] step 38660 | loss=0.4685 pred=0.0466 kl=1.2434 recon=0.2976 | grad=1.43 | 47.6s
Epoch 358/1000 | avg_loss=0.4662 pred=0.0467 kl=1.2428 recon=0.2953 | 49.5s
  New best model (loss=0.4662)
  [359/1000] step 38680 | loss=0.4682 pred=0.0468 kl=1.2408 recon=0.2973 | grad=1.59 | 7.5s
  [359/1000] step 38700 | loss=0.4667 pred=0.0467 kl=1.2411 recon=0.2959 | grad=1.70 | 16.7s
  [359/1000] step 38720 | loss=0.4666 pred=0.0467 kl=1.2439 recon=0.2956 | grad=1.47 | 25.8s
  [359/1000] step 38740 | loss=0.4686 pred=0.0468 kl=1.2421 recon=0.2976 | grad=1.25 | 35.0s
  [359/1000] step 38760 | loss=0.4674 pred=0.0469 kl=1.2427 recon=0.2963 | grad=1.44 | 43.3s
Epoch 359/1000 | avg_loss=0.4670 pred=0.0467 kl=1.2422 recon=0.2961 | 49.1s
  [360/1000] step 38780 | loss=0.4666 pred=0.0468 kl=1.2425 recon=0.2956 | grad=1.00 | 3.9s
  [360/1000] step 38800 | loss=0.4636 pred=0.0466 kl=1.2426 recon=0.2927 | grad=0.36 | 13.4s
  [360/1000] step 38820 | loss=0.4640 pred=0.0466 kl=1.2434 recon=0.2930 | grad=0.45 | 22.9s
  [360/1000] step 38840 | loss=0.4653 pred=0.0466 kl=1.2439 recon=0.2943 | grad=1.06 | 32.2s
  [360/1000] step 38860 | loss=0.4636 pred=0.0463 kl=1.2430 recon=0.2930 | grad=0.89 | 41.6s
  [360/1000] step 38880 | loss=0.4675 pred=0.0468 kl=1.2424 recon=0.2965 | grad=1.94 | 51.2s
Epoch 360/1000 | avg_loss=0.4660 pred=0.0466 kl=1.2430 recon=0.2951 | 51.3s
  New best model (loss=0.4660)
  [361/1000] step 38900 | loss=0.4666 pred=0.0466 kl=1.2426 recon=0.2958 | grad=1.57 | 9.5s
  [361/1000] step 38920 | loss=0.4650 pred=0.0466 kl=1.2428 recon=0.2941 | grad=1.58 | 18.6s
  [361/1000] step 38940 | loss=0.4637 pred=0.0466 kl=1.2425 recon=0.2929 | grad=0.76 | 27.8s
  [361/1000] step 38960 | loss=0.4640 pred=0.0467 kl=1.2423 recon=0.2931 | grad=0.41 | 35.9s
  [361/1000] step 38980 | loss=0.4672 pred=0.0466 kl=1.2443 recon=0.2962 | grad=0.43 | 45.1s
Epoch 361/1000 | avg_loss=0.4660 pred=0.0467 kl=1.2428 recon=0.2950 | 48.9s
  New best model (loss=0.4660)
  [362/1000] step 39000 | loss=0.4679 pred=0.0467 kl=1.2431 recon=0.2969 | grad=0.88 | 5.7s
  [362/1000] step 39020 | loss=0.4624 pred=0.0464 kl=1.2435 recon=0.2916 | grad=0.58 | 14.7s
  [362/1000] step 39040 | loss=0.4665 pred=0.0465 kl=1.2445 recon=0.2956 | grad=0.50 | 24.0s
  [362/1000] step 39060 | loss=0.4702 pred=0.0467 kl=1.2404 recon=0.2994 | grad=1.71 | 33.3s
  [362/1000] step 39080 | loss=0.4643 pred=0.0467 kl=1.2422 recon=0.2934 | grad=1.02 | 42.6s
Epoch 362/1000 | avg_loss=0.4664 pred=0.0466 kl=1.2431 recon=0.2954 | 50.1s
  [363/1000] step 39100 | loss=0.4676 pred=0.0466 kl=1.2445 recon=0.2966 | grad=0.47 | 2.2s
  [363/1000] step 39120 | loss=0.4643 pred=0.0465 kl=1.2442 recon=0.2934 | grad=0.56 | 11.7s
  [363/1000] step 39140 | loss=0.4671 pred=0.0465 kl=1.2440 recon=0.2962 | grad=0.40 | 21.1s
  [363/1000] step 39160 | loss=0.4650 pred=0.0464 kl=1.2450 recon=0.2942 | grad=1.37 | 29.6s
  [363/1000] step 39180 | loss=0.4645 pred=0.0465 kl=1.2443 recon=0.2936 | grad=0.38 | 38.9s
  [363/1000] step 39200 | loss=0.4634 pred=0.0465 kl=1.2448 recon=0.2924 | grad=0.60 | 48.2s
Epoch 363/1000 | avg_loss=0.4649 pred=0.0466 kl=1.2437 recon=0.2939 | 50.0s
  New best model (loss=0.4649)
  [364/1000] step 39220 | loss=0.4678 pred=0.0467 kl=1.2468 recon=0.2963 | grad=1.37 | 7.7s
  [364/1000] step 39240 | loss=0.4738 pred=0.0468 kl=1.2434 recon=0.3027 | grad=1.91 | 17.7s
  [364/1000] step 39260 | loss=0.4639 pred=0.0467 kl=1.2401 recon=0.2932 | grad=1.66 | 27.5s
  [364/1000] step 39280 | loss=0.4667 pred=0.0468 kl=1.2415 recon=0.2958 | grad=1.40 | 36.7s
  [364/1000] step 39300 | loss=0.4634 pred=0.0466 kl=1.2414 recon=0.2927 | grad=1.17 | 45.9s
Epoch 364/1000 | avg_loss=0.4660 pred=0.0467 kl=1.2430 recon=0.2950 | 51.4s
  [365/1000] step 39320 | loss=0.4669 pred=0.0468 kl=1.2426 recon=0.2959 | grad=1.07 | 3.8s


Losses look fine — not diverging. The problem isn't overfitting. Let me try inference with an earlier checkpoint to confirm:


python3 inference.py --checkpoint output/checkpoints/epoch_0070.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v1_e070_vocos_synth.wav --max_steps 500

Test earlier epoch 70 checkpoint synthesis


python3 -c "
import torch, numpy as np, torchaudio

# Compare norms of synthesized outputs
for f in ['output/v1_e070_vocos_synth.wav', 'output/v1_e350_vocos_synth.wav', 'output/test_ar_synth.wav']:
    try:
        w, sr = torchaudio.load(f)
        print(f'{f}: duration={w.shape[1]/sr:.2f}s, rms={w.pow(2).mean().sqrt():.4f}, max={w.abs().max():.4f}')
    except: pass
"

Compare audio stats across checkpoints


Model loaded from output/checkpoints/epoch_0070.pt
Saved: output/v1_e070_vocos_synth.wav (7.41s)


output/v1_e070_vocos_synth.wav: duration=7.41s, rms=0.1560, max=0.8025
output/v1_e350_vocos_synth.wav: duration=2.34s, rms=0.0262, max=0.7986
output/test_ar_synth.wav: duration=8.57s, rms=0.0231, max=0.3219


There it is. **Epoch 70 gives 7.41s output, epoch 350 only 2.34s**. The model is collapsing at later epochs — the prediction loss is so low (0.046) that embeddings become too uniform and the AR loop's norm-based stopping criterion triggers too early.

Let me also try the `best.pt` and run v2 with its earliest checkpoint:


python3 inference.py --checkpoint output/checkpoints/epoch_0100.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v1_e100_vocos_synth.wav --max_steps 500

Test epoch 100 synthesis


cd /home/ubuntu/lewm-tts/v2 && python3 inference.py --checkpoint output/checkpoints/epoch_0010.pt --text "भारत एक महान देश है और हम सब भारतवासी हैं।" --mode synthesize --output output/v2_e010_synth.wav --max_steps 500

Test v2 early checkpoint synthesis


Model loaded from output/checkpoints/epoch_0100.pt
Saved: output/v1_e100_vocos_synth.wav (17.99s)


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loading DAC 24kHz...
Model loaded. Start scale: 137.140
Saved: output/v2_e010_synth.wav (6.68s)


Confirmed. Here's the pattern:

| Checkpoint | Duration | Quality |
|-----------|----------|---------|
| **v1 epoch 70** | **7.41s** | Good |
| **v1 epoch 100** | **17.99s** | Good (maybe too long) |
| v1 epoch 350 | 2.34s | Collapsed — stops too early |
| v2 epoch 10 | 6.68s | Better than v2 epoch 150 |
| v2 epoch 150 | 0.55s | Collapsed |

The problem is clear: **later epochs overfit** — prediction loss gets so low that all embeddings become nearly identical, so the norm-based stopping criterion (`std < 0.01 * mean`) fires immediately. Earlier checkpoints (epoch 70-100 for v1, epoch 10 for v2) produce better AR synthesis.

The best checkpoints for synthesis are **v1 epoch 70-100** and **v2 epoch 10**. Give those a listen!


but everuuthing pure noise .. u dng infernce some wrong


The user says everything is pure noise. Let me look more carefully at how inference was done previously when it worked. Let me check the old working wav files and compare. Maybe the issue is with the Vocos padding (80→100 mels) or something else in the inference pipeline.

Let me look at the old test_sample_epoch10.wav - that was generated at epoch 10 and was 613KB. How was it generated? Maybe it was generated during training, not through inference.py.

Let me check train.py for any sample generation code.


Let me check how the old working samples were actually generated — maybe it was done during training, not through inference.py.


grep -n "sample\|synth\|vocos\|wav\|save_audio\|torchaudio.save" /home/ubuntu/lewm-tts/train.py | head -30

Check how training generates samples


grep -n "vocos\|sample\|synth\|wav" /home/ubuntu/lewm-tts/output/train_v1_resumed.log | head -20

Check training log for sample generation


6:Dataset: 6959 samples loaded
7:Training: 6959 samples, 108 batches/epoch


43:    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")


Training doesn't generate samples. The old working wavs must have been generated differently. Let me look at what `test_sample_epoch10.wav` actually is — and check bash history:


cat /home/ubuntu/.bash_history 2>/dev/null | grep -i "infer\|synth\|vocos\|test_sample\|reconstruct" | tail -30

Check bash history for old inference commands


Check current inference.py state


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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        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.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        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, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        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):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


git clone https://github.com/MayaResearch/veena3_inference.git
git clone https://github.com/MayaResearch/veena3_inference.git .
mkdir soprano_infernce
cd soprano_infernce/
cd kani-tts-inference/


The issue is the Vocos padding — padding 80 mels to 100 with zeros is wrong. The Vocos mel-24khz model expects 100-bin mel spectrograms with specific frequency binning. Just zero-padding the upper 20 bins produces garbage.

Let me check what the old working vocos files used. The `vocos_groundtruth.wav` was generated before — how? Let me check if there's a different script, or if Vocos was used differently.

Actually, looking at the Vocos API more carefully — `Vocos.from_pretrained("charactr/vocos-mel-24khz")` has a `forward` method that takes raw audio, and `decode` takes mel features.

The core problem is that our model outputs 80-bin mels but Vocos expects 100, and zero-padding introduces artifacts. I need to figure out how the previous vocos files were actually generated — maybe the mel extraction was done differently, or a different Vocos model was used entirely. I'm considering whether the model was originally designed to output 100-bin mels for Vocos compatibility, or if there's a different vocoder in play—let me check the model config to see what n_mels is actually set to.


The Vocos padding (80→100 mels with zeros) is definitely wrong — that'll produce noise. Let me check the actual model config and how it was originally used.


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
ckpt = torch.load('output/checkpoints/epoch_0070.pt', map_location='cpu', weights_only=False)
print('Config:', ckpt['config'])
"

Check model config for n_mels


Config: {'d_model': 256, 'nhead': 4, 'n_mels': 80, 'text_vocab_size': 256, 'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6, 'dropout': 0.1, 'kl_weight': 0.1, 'recon_weight': 1.0}


     1→"""
     2→LeWM TTS Model — JEPA-based Text-to-Speech
     3→
     4→Architecture:
     5→  - TextEncoder: Character-level transformer for Hindi/Devanagari text
     6→  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
     7→  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
     8→  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
     9→"""
    10→
    11→import math
    12→import torch
    13→import torch.nn as nn
    14→import torch.nn.functional as F
    15→
    16→
    17→# ─── Text Encoder ───────────────────────────────────────────────────────────
    18→
    19→class TextEncoder(nn.Module):
    20→    """Character-level transformer encoder for Hindi text."""
    21→
    22→    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
    23→        super().__init__()
    24→        self.d_model = d_model
    25→
    26→        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
    27→        self.char_embed = nn.Embedding(vocab_size, d_model)
    28→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
    29→        self.dropout = nn.Dropout(dropout)
    30→
    31→        encoder_layer = nn.TransformerEncoderLayer(
    32→            d_model=d_model, nhead=nhead,
    33→            dim_feedforward=d_model * 4,
    34→            dropout=dropout, batch_first=True,
    35→            activation="gelu",
    36→        )
    37→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    38→        self.proj = nn.Linear(d_model, d_model)
    39→
    40→    def forward(self, text_tokens, text_mask=None):
    41→        """
    42→        Args:
    43→            text_tokens: [B, T_text] long tensor of byte-level token IDs
    44→            text_mask: [B, T_text] bool, True = padding
    45→        Returns:
    46→            text_emb: [B, T_text, d_model]
    47→        """
    48→        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
    49→        x = self.pos_embed(x)
    50→        x = self.dropout(x)
    51→        x = self.transformer(x, src_key_padding_mask=text_mask)
    52→        x = self.proj(x)
    53→        return x
    54→
    55→
    56→# ─── Audio Encoder ──────────────────────────────────────────────────────────
    57→
    58→class AudioEncoder(nn.Module):
    59→    """
    60→    Encodes mel spectrograms into a sequence of audio embeddings.
    61→    Uses 1D CNN for downsampling + Transformer for contextualization.
    62→    """
    63→
    64→    def __init__(self, n_mels=80, d_model=256, nhead=4, num_layers=4,
    65→                 downsample_factor=4, dropout=0.1):
    66→        super().__init__()
    67→        self.d_model = d_model
    68→        self.downsample_factor = downsample_factor
    69→
    70→        # CNN downsampling: reduces temporal resolution by downsample_factor
    71→        self.conv_pre = nn.Sequential(
    72→            nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    73→            nn.GELU(),
    74→            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    75→            nn.GELU(),
    76→            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    77→            nn.GELU(),
    78→        )  # Total stride = 4
    79→
    80→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
    81→        self.dropout = nn.Dropout(dropout)
    82→
    83→        encoder_layer = nn.TransformerEncoderLayer(
    84→            d_model=d_model, nhead=nhead,
    85→            dim_feedforward=d_model * 4,
    86→            dropout=dropout, batch_first=True,
    87→            activation="gelu",
    88→        )
    89→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    90→
    91→        # Project to embedding with mu and logvar for Gaussian regularizer
    92→        self.proj_mu = nn.Linear(d_model, d_model)
    93→        self.proj_logvar = nn.Linear(d_model, d_model)
    94→
    95→    def forward(self, mel, mel_mask=None):
    96→        """
    97→        Args:
    98→            mel: [B, n_mels, T_mel] — log mel spectrogram
    99→            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
   100→        Returns:
   101→            z: [B, T_down, d_model] — sampled latent embeddings
   102→            mu: [B, T_down, d_model]
   103→            logvar: [B, T_down, d_model]
   104→        """
   105→        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
   106→        x = self.conv_pre(mel)
   107→        x = x.transpose(1, 2)  # [B, T_down, d_model]
   108→
   109→        # Downsample mask if provided
   110→        if mel_mask is not None:
   111→            # Approximate downsampled mask
   112→            T_down = x.shape[1]
   113→            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
   114→
   115→        x = self.pos_embed(x)
   116→        x = self.dropout(x)
   117→        x = self.transformer(x, src_key_padding_mask=mel_mask)
   118→
   119→        mu = self.proj_mu(x)
   120→        logvar = self.proj_logvar(x)
   121→
   122→        # Reparameterization trick
   123→        if self.training:
   124→            std = torch.exp(0.5 * logvar)
   125→            eps = torch.randn_like(std)
   126→            z = mu + eps * std
   127→        else:
   128→            z = mu
   129→
   130→        return z, mu, logvar
   131→
   132→
   133→# ─── JEPA Predictor ─────────────────────────────────────────────────────────
   134→
   135→class JEPAPredictor(nn.Module):
   136→    """
   137→    Predicts next audio embedding given:
   138→    - Text context (from TextEncoder)
   139→    - Previous audio embeddings (from AudioEncoder)
   140→
   141→    Uses cross-attention to condition on text, and causal self-attention
   142→    over the audio embedding sequence.
   143→    """
   144→
   145→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
   146→        super().__init__()
   147→        self.d_model = d_model
   148→
   149→        self.audio_input_proj = nn.Linear(d_model, d_model)
   150→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
   151→        self.dropout = nn.Dropout(dropout)
   152→
   153→        # Transformer decoder with causal self-attention + cross-attention to text
   154→        decoder_layer = nn.TransformerDecoderLayer(
   155→            d_model=d_model, nhead=nhead,
   156→            dim_feedforward=d_model * 4,
   157→            dropout=dropout, batch_first=True,
   158→            activation="gelu",
   159→        )
   160→        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
   161→
   162→        # Output projection to predict next embedding
   163→        self.output_proj = nn.Sequential(
   164→            nn.Linear(d_model, d_model),
   165→            nn.GELU(),
   166→            nn.Linear(d_model, d_model),
   167→        )
   168→
   169→    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
   170→        """
   171→        Args:
   172→            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
   173→            text_emb: [B, T_text, d_model] — text embeddings
   174→            audio_mask: [B, T_audio] bool, True = padding
   175→            text_mask: [B, T_text] bool, True = padding
   176→        Returns:
   177→            predicted: [B, T_audio, d_model] — predicted next embeddings
   178→        """
   179→        x = self.audio_input_proj(audio_emb)
   180→        x = self.pos_embed(x)
   181→        x = self.dropout(x)
   182→
   183→        # Causal mask: each position can only attend to previous positions
   184→        T = x.shape[1]
   185→        causal_mask = torch.triu(
   186→            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
   187→        )
   188→
   189→        predicted = self.transformer(
   190→            x, text_emb,
   191→            tgt_mask=causal_mask,
   192→            tgt_key_padding_mask=audio_mask,
   193→            memory_key_padding_mask=text_mask,
   194→        )
   195→
   196→        predicted = self.output_proj(predicted)
   197→        return predicted
   198→
   199→
   200→# ─── Positional Encoding ────────────────────────────────────────────────────
   201→
   202→class SinusoidalPositionalEncoding(nn.Module):
   203→    def __init__(self, d_model, max_len=8192):
   204→        super().__init__()
   205→        pe = torch.zeros(max_len, d_model)
   206→        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
   207→        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
   208→        pe[:, 0::2] = torch.sin(position * div_term)
   209→        pe[:, 1::2] = torch.cos(position * div_term)
   210→        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
   211→
   212→    def forward(self, x):
   213→        return x + self.pe[:, :x.shape[1]]
   214→
   215→
   216→# ─── Mel Decoder ────────────────────────────────────────────────────────────
   217→
   218→class MelDecoder(nn.Module):
   219→    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""
   220→
   221→    def __init__(self, d_model=256, n_mels=80):
   222→        super().__init__()
   223→        self.proj = nn.Sequential(
   224→            nn.Linear(d_model, d_model),
   225→            nn.GELU(),
   226→        )
   227→        self.upsample = nn.Sequential(
   228→            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
   229→            nn.GELU(),
   230→            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
   231→            nn.GELU(),
   232→            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),
   233→        )
   234→
   235→    def forward(self, z):
   236→        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
   237→        x = self.proj(z)
   238→        x = x.transpose(1, 2)
   239→        return self.upsample(x)
   240→
   241→
   242→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   243→
   244→class LeWMTTS(nn.Module):
   245→    """
   246→    Complete LeWM TTS model.
   247→
   248→    Training: text + mel → audio embeddings → predict next embedding
   249→    Two losses:
   250→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   251→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   252→    """
   253→
   254→    def __init__(self, config):
   255→        super().__init__()
   256→        d_model = config.get("d_model", 256)
   257→        nhead = config.get("nhead", 4)
   258→        n_mels = config.get("n_mels", 80)
   259→        text_vocab = config.get("text_vocab_size", 256)
   260→        text_layers = config.get("text_encoder_layers", 4)
   261→        audio_layers = config.get("audio_encoder_layers", 4)
   262→        predictor_layers = config.get("predictor_layers", 6)
   263→        dropout = config.get("dropout", 0.1)
   264→        self.kl_weight = config.get("kl_weight", 0.01)
   265→
   266→        self.text_encoder = TextEncoder(
   267→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   268→            num_layers=text_layers, dropout=dropout,
   269→        )
   270→        self.audio_encoder = AudioEncoder(
   271→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   272→            num_layers=audio_layers, dropout=dropout,
   273→        )
   274→        self.predictor = JEPAPredictor(
   275→            d_model=d_model, nhead=nhead,
   276→            num_layers=predictor_layers, dropout=dropout,
   277→        )
   278→
   279→        # Mel decoder: reconstructs mel from latent embeddings
   280→        # Needed so the latent space is grounded to actual audio
   281→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
   282→        self.recon_weight = config.get("recon_weight", 1.0)
   283→
   284→        self.config = config
   285→
   286→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
   287→        """
   288→        Args:
   289→            mel: [B, n_mels, T_mel]
   290→            text_tokens: [B, T_text]
   291→            mel_mask: [B, T_mel] bool
   292→            text_mask: [B, T_text] bool
   293→        Returns:
   294→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   295→        """
   296→        # Encode text
   297→        text_emb = self.text_encoder(text_tokens, text_mask)
   298→
   299→        # Encode audio → latent embeddings
   300→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   301→
   302→        # Predict next embedding: input z[:-1], target z[1:]
   303→        input_emb = z[:, :-1]     # [B, T-1, d]
   304→        target_emb = z[:, 1:]     # [B, T-1, d]
   305→
   306→        # Adjust mask for downsampled sequence
   307→        if mel_mask is not None:
   308→            T_down = z.shape[1]
   309→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   310→            pred_mask = ds_mask[:, :-1]
   311→        else:
   312→            pred_mask = None
   313→
   314→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   315→
   316→        # ─── Loss 1: Next-embedding prediction (MSE) ───
   317→        # Mask out padding positions
   318→        if pred_mask is not None:
   319→            valid = ~pred_mask  # True = valid
   320→            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
   321→            prediction_loss = F.mse_loss(
   322→                predicted * valid, target_emb * valid, reduction="sum"
   323→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   324→        else:
   325→            prediction_loss = F.mse_loss(predicted, target_emb)
   326→
   327→        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
   328→        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
   329→        if mel_mask is not None:
   330→            T_down = mu.shape[1]
   331→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   332→            valid = (~ds_mask).unsqueeze(-1).float()
   333→            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
   334→            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
   335→        else:
   336→            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
   337→
   338→        # ─── Loss 3: Mel reconstruction ───
   339→        # Decode latents back to mel and compare with input
   340→        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
   341→        # Trim/pad to match input mel length
   342→        T_mel = mel.shape[2]
   343→        T_recon = mel_recon.shape[2]
   344→        if T_recon > T_mel:
   345→            mel_recon = mel_recon[:, :, :T_mel]
   346→        elif T_recon < T_mel:
   347→            mel = mel[:, :, :T_recon]
   348→            if mel_mask is not None:
   349→                mel_mask = mel_mask[:, :T_recon]
   350→
   351→        if mel_mask is not None:
   352→            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
   353→            recon_loss = F.l1_loss(
   354→                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
   355→            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
   356→        else:
   357→            recon_loss = F.l1_loss(mel_recon, mel)
   358→
   359→        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss
   360→
   361→        return {
   362→            "total_loss": total_loss,
   363→            "prediction_loss": prediction_loss,
   364→            "kl_loss": kl_loss,
   365→            "recon_loss": recon_loss,
   366→        }
   367→
   368→    def encode_audio(self, mel):
   369→        """Encode mel to latent embeddings (for inference)."""
   370→        z, mu, _ = self.audio_encoder(mel)
   371→        return mu  # Use mean at inference
   372→
   373→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   374→        """Predict next audio embedding autoregressively."""
   375→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   376→        return predicted[:, -1:]  # Only the last prediction
   377→
   378→
   379→def count_parameters(model):
   380→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   381→
   382→
   383→def build_model(config=None):
   384→    if config is None:
   385→        config = {
   386→            "d_model": 256,
   387→            "nhead": 4,
   388→            "n_mels": 80,
   389→            "text_vocab_size": 256,
   390→            "text_encoder_layers": 4,
   391→            "audio_encoder_layers": 4,
   392→            "predictor_layers": 6,
   393→            "dropout": 0.1,
   394→            "kl_weight": 0.01,
   395→        }
   396→    model = LeWMTTS(config)
   397→    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
   398→    return model, config
   399→

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# ─── Text Encoder ───────────────────────────────────────────────────────────

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=80, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        self.conv_pre = nn.Sequential(
            nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
            nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
        )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

# ─── JEPA Predictor ─────────────────────────────────────────────────────────

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

    def forward(self, x):
        return x + self.pe[:, :x.shape[1]]

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""

    def __init__(self, d_model=256, n_mels=80):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)
        return self.upsample(x)

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 80)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
        self.recon_weight = config.get("recon_weight", 1.0)

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 80,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


OK I see the problem now. Padding 80→100 mels with zeros is garbage — Vocos was trained on 100-bin mels with specific frequency bins. Let me check how Vocos actually works for 80-mel input:


python3 -c "
from vocos import Vocos
import torch

v = Vocos.from_pretrained('charactr/vocos-mel-24khz')

# Check the feature extractor mel config
fe = v.feature_extractor
print(dir(fe.mel_spec))
ms = fe.mel_spec
print('n_fft:', ms.spectrogram.n_fft)
print('hop_length:', ms.spectrogram.hop_length)
print('n_mels:', ms.mel_scale.n_mels)
print('sample_rate:', ms.sample_rate)
"

Check Vocos mel config


['T_destination', '__annotations__', '__call__', '__class__', '__constants__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattr__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setstate__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_apply', '_backward_hooks', '_backward_pre_hooks', '_buffers', '_call_impl', '_compiled_call_impl', '_forward_hooks', '_forward_hooks_always_called', '_forward_hooks_with_kwargs', '_forward_pre_hooks', '_forward_pre_hooks_with_kwargs', '_get_backward_hooks', '_get_backward_pre_hooks', '_get_name', '_is_full_backward_hook', '_load_from_state_dict', '_load_state_dict_post_hooks', '_load_state_dict_pre_hooks', '_maybe_warn_non_full_backward_hook', '_modules', '_named_members', '_non_persistent_buffers_set', '_parameters', '_register_load_state_dict_pre_hook', '_register_state_dict_hook', '_replicate_for_data_parallel', '_save_to_state_dict', '_slow_forward', '_state_dict_hooks', '_state_dict_pre_hooks', '_version', '_wrapped_call_impl', 'add_module', 'apply', 'bfloat16', 'buffers', 'call_super_init', 'children', 'compile', 'cpu', 'cuda', 'double', 'dump_patches', 'eval', 'extra_repr', 'f_max', 'f_min', 'float', 'forward', 'get_buffer', 'get_extra_state', 'get_parameter', 'get_submodule', 'half', 'hop_length', 'ipu', 'load_state_dict', 'mel_scale', 'modules', 'mtia', 'n_fft', 'n_mels', 'named_buffers', 'named_children', 'named_modules', 'named_parameters', 'normalized', 'pad', 'parameters', 'power', 'register_backward_hook', 'register_buffer', 'register_forward_hook', 'register_forward_pre_hook', 'register_full_backward_hook', 'register_full_backward_pre_hook', 'register_load_state_dict_post_hook', 'register_load_state_dict_pre_hook', 'register_module', 'register_parameter', 'register_state_dict_post_hook', 'register_state_dict_pre_hook', 'requires_grad_', 'sample_rate', 'set_extra_state', 'set_submodule', 'share_memory', 'spectrogram', 'state_dict', 'to', 'to_empty', 'train', 'training', 'type', 'win_length', 'xpu', 'zero_grad']
n_fft: 1024
hop_length: 256
n_mels: 100
sample_rate: 24000


Vocos uses n_fft=1024, hop=256, sr=24000 — same as our model, but **100 mels instead of 80**. The correct fix is to recompute the mel with 100 bins from the waveform. But we can't do that because we need the vocoder to GET the waveform.

The right approach: use Griffin-Lim to get a rough waveform first, then use Vocos's own pipeline (audio → 100-mel → decode) to refine it. OR just use Griffin-Lim properly since that's what actually works with our 80-mel model.

Let me revert to Griffin-Lim and make sure it's correct:


Edit

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

        # Vocoder: Griffin-Lim → Vocos refinement
        self.gl_vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocoder_decode(self, mel):
        """Decode mel (B, 80, T) → waveform. Griffin-Lim → Vocos refinement."""
        # Step 1: Griffin-Lim to get rough waveform from 80-mel
        waveform = self.gl_vocoder(mel)
        # Step 2: Vocos refinement — extract 100-mel and decode
        waveform = waveform.unsqueeze(0)  # [1, 1, T]
        waveform = self._vocos(waveform)  # Vocos forward: audio → mel(100) → decode
        return waveform


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

        # Vocoder: Griffin-Lim → Vocos refinement
        self.gl_vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocoder_decode(self, mel):
        """Decode mel (B, 80, T) → waveform. Griffin-Lim → Vocos refinement."""
        # Step 1: Griffin-Lim to get rough waveform from 80-mel
        waveform = self.gl_vocoder(mel)
        # Step 2: Vocos refinement — extract 100-mel and decode
        waveform = waveform.unsqueeze(0)  # [1, 1, T]
        waveform = self._vocos(waveform)  # Vocos forward: audio → mel(100) → decode
        return waveform

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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()

        # Vocoder (Vocos)
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocos_decode(self, mel):
        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""
        import torch.nn.functional as F
        # Pad 80 → 100 mels
        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)
        return self._vocos.decode(mel_padded)

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embeddings, dim=1)
            next_emb = self.model.predict_next(context, text_emb)

            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Stopping: track norm stability — stop if converged
            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)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        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.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        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, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        # Decode
        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):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

+        # Vocoder: Griffin-Lim → Vocos refinement

+        self.gl_vocoder = GriffinLimVocoder(

+            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,

-        """Decode mel (B, 80, T) → waveform via Vocos (expects 100 mels)."""

-        import torch.nn.functional as F

-        mel_padded = F.pad(mel, (0, 0, 0, 20), mode='constant', value=0)

-        return self._vocos.decode(mel_padded)

+        """Decode mel (B, 80, T) → waveform. Griffin-Lim → Vocos refinement."""

+        # Step 1: Griffin-Lim to get rough waveform from 80-mel

+        waveform = self.gl_vocoder(mel)

+        # Step 2: Vocos refinement — extract 100-mel and decode

+        waveform = waveform.unsqueeze(0)  # [1, 1, T]

+        waveform = self._vocos(waveform)  # Vocos forward: audio → mel(100) → decode