hi


Hi! How can I help you today?


'/home/ubuntu/lewm-tts' run infernce on 2 versions


[Request interrupted by user]


run infernce on 2 version latets checkpoints '/home/ubuntu/lewm-tts'


The user wants to run inference on the 2 latest checkpoints from their LeWM TTS project. Let me first look at the project structure and find the checkpoints and inference script.


Let me look at the project structure, checkpoints, and inference code.


ToolSearch

Question: select:Bash,Glob,Read,Agent


Question: select:Bash,Glob,Read,Agent


Let me find the checkpoints and inference script in the project.


Bash

ls /home/ubuntu/lewm-tts/output/checkpoints/ 2>/dev/null || ls /home/ubuntu/lewm-tts/output/ 2>/dev/null | head -30

List checkpoints in output directory


/home/ubuntu/lewm-tts/inference.py
/home/ubuntu/lewm-tts/__pycache__/inference.cpython-310.pyc
/home/ubuntu/lewm-tts/v2/inference.py
/home/ubuntu/lewm-tts/v2/__pycache__/inference.cpython-310.pyc


best.pt
epoch_0070.pt
epoch_0100.pt
epoch_0150.pt
epoch_0200.pt
epoch_0250.pt
epoch_0300.pt
epoch_0350.pt


The 2 latest checkpoints are `epoch_0350.pt` and `epoch_0300.pt`. Let me check the inference scripts to understand how to run them.


     1→"""
     2→Inference pipeline for LeWM TTS.
     3→Text → JEPA predictor → mel reconstruction → vocoder → waveform
     4→"""
     5→
     6→import torch
     7→import torch.nn as nn
     8→import torchaudio
     9→import json
    10→import argparse
    11→import numpy as np
    12→from pathlib import Path
    13→
    14→from model import LeWMTTS
    15→
    16→
    17→class GriffinLimVocoder:
    18→    """Griffin-Lim vocoder as fallback."""
    19→
    20→    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
    21→        self.n_fft = n_fft
    22→        self.hop_length = hop_length
    23→        self.sample_rate = sample_rate
    24→
    25→        self.inverse_mel = torchaudio.transforms.InverseMelScale(
    26→            n_stft=n_fft // 2 + 1,
    27→            n_mels=n_mels,
    28→            sample_rate=sample_rate,
    29→        )
    30→        self.griffin_lim = torchaudio.transforms.GriffinLim(
    31→            n_fft=n_fft,
    32→            hop_length=hop_length,
    33→            n_iter=64,
    34→            power=1.0,
    35→        )
    36→
    37→    def __call__(self, mel):
    38→        mel = torch.exp(mel)
    39→        spec = self.inverse_mel(mel)
    40→        waveform = self.griffin_lim(spec)
    41→        return waveform
    42→
    43→
    44→class LeWMTTSInference:
    45→    """Full inference pipeline."""
    46→
    47→    def __init__(self, checkpoint_path, device="cuda"):
    48→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    49→
    50→        # Load checkpoint
    51→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    52→        config = ckpt["config"]
    53→
    54→        # Build model
    55→        self.model = LeWMTTS(config).to(self.device)
    56→        self.model.load_state_dict(ckpt["model"])
    57→        self.model.eval()
    58→
    59→        # Vocoder
    60→        self.vocoder = GriffinLimVocoder(
    61→            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
    62→        )
    63→
    64→        self.config = config
    65→
    66→        # Compute start token from training data statistics
    67→        # Use a random normal scaled to match training embedding norms
    68→        self._start_emb_scale = 5.0  # approximate from diagnostics
    69→
    70→        print(f"Model loaded from {checkpoint_path}")
    71→
    72→    def text_to_tokens(self, text):
    73→        tokens = list(text.encode("utf-8"))
    74→        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
    75→
    76→    @torch.no_grad()
    77→    def synthesize(self, text, max_steps=300, temperature=0.0):
    78→        """
    79→        Generate speech from text using AR prediction.
    80→        """
    81→        text_tokens = self.text_to_tokens(text)
    82→        text_emb = self.model.text_encoder(text_tokens)
    83→
    84→        d = self.config["d_model"]
    85→
    86→        # Start with a Gaussian sample matching training distribution
    87→        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
    88→        all_embeddings = [start_emb]
    89→
    90→        prev_norms = []
    91→
    92→        for step in range(max_steps):
    93→            context = torch.cat(all_embeddings, dim=1)
    94→            next_emb = self.model.predict_next(context, text_emb)
    95→
    96→            if temperature > 0:
    97→                noise = torch.randn_like(next_emb) * temperature
    98→                next_emb = next_emb + noise
    99→
   100→            all_embeddings.append(next_emb)
   101→
   102→            # Stopping: track norm stability — stop if converged
   103→            norm = next_emb.norm().item()
   104→            prev_norms.append(norm)
   105→            if len(prev_norms) > 20:
   106→                recent = prev_norms[-20:]
   107→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
   108→                    break
   109→
   110→        audio_embs = torch.cat(all_embeddings, dim=1)
   111→
   112→        # Decode to mel
   113→        mel = self.model.mel_decoder(audio_embs)
   114→
   115→        # Vocoder
   116→        waveform = self.vocoder(mel.cpu())
   117→        return waveform.squeeze().numpy(), 24000
   118→
   119→    @torch.no_grad()
   120→    def reconstruct(self, mel_path):
   121→        """
   122→        Teacher-forced reconstruction: encode real mel → decode back.
   123→        Tests mel decoder + vocoder quality without AR drift.
   124→        """
   125→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
   126→        # [1, n_mels, T]
   127→
   128→        # Encode
   129→        z, mu, logvar = self.model.audio_encoder(mel)
   130→
   131→        # Decode from mean embeddings
   132→        mel_recon = self.model.mel_decoder(mu)
   133→
   134→        # Vocoder
   135→        waveform = self.vocoder(mel_recon.cpu())
   136→        return waveform.squeeze().numpy(), 24000
   137→
   138→    @torch.no_grad()
   139→    def reconstruct_with_prediction(self, mel_path):
   140→        """
   141→        Hybrid: encode real mel, run predictor on real embeddings,
   142→        decode predicted embeddings. Tests prediction quality without AR drift.
   143→        """
   144→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
   145→
   146→        # Encode audio
   147→        z, mu, logvar = self.model.audio_encoder(mel)
   148→
   149→        # We need text too — use dummy text (empty)
   150→        # This tests if the predictor can reproduce the audio embeddings
   151→        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
   152→        text_emb = self.model.text_encoder(text_tokens)
   153→
   154→        # Run predictor on real embeddings (teacher-forced)
   155→        predicted = self.model.predictor(mu[:, :-1], text_emb)
   156→
   157→        # Use real first embedding + predicted rest
   158→        combined = torch.cat([mu[:, :1], predicted], dim=1)
   159→
   160→        # Decode
   161→        mel_recon = self.model.mel_decoder(combined)
   162→        waveform = self.vocoder(mel_recon.cpu())
   163→        return waveform.squeeze().numpy(), 24000
   164→
   165→    def save_audio(self, waveform, sr, output_path):
   166→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
   167→        torchaudio.save(output_path, wav_tensor, sr)
   168→        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
   169→
   170→
   171→def main():
   172→    parser = argparse.ArgumentParser()
   173→    parser.add_argument("--checkpoint", required=True)
   174→    parser.add_argument("--text", default=None, help="Text to synthesize")
   175→    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
   176→    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
   177→                        default="synthesize")
   178→    parser.add_argument("--output", default="output.wav")
   179→    parser.add_argument("--max_steps", type=int, default=300)
   180→    parser.add_argument("--temperature", type=float, default=0.0)
   181→    args = parser.parse_args()
   182→
   183→    tts = LeWMTTSInference(args.checkpoint)
   184→
   185→    if args.mode == "synthesize":
   186→        assert args.text, "Need --text for synthesis"
   187→        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
   188→    elif args.mode == "reconstruct":
   189→        assert args.mel, "Need --mel for reconstruction"
   190→        waveform, sr = tts.reconstruct(args.mel)
   191→    elif args.mode == "predict":
   192→        assert args.mel, "Need --mel for prediction test"
   193→        waveform, sr = tts.reconstruct_with_prediction(args.mel)
   194→
   195→    tts.save_audio(waveform, sr, args.output)
   196→
   197→
   198→if __name__ == "__main__":
   199→    main()
   200→

"""
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
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

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


     1→"""Inference for LeWM TTS v2 (DAC-based)."""
     2→
     3→import torch
     4→import torchaudio
     5→import dac
     6→import numpy as np
     7→import argparse
     8→from model import LeWMTTS
     9→
    10→
    11→class LeWMTTSInference:
    12→    def __init__(self, checkpoint_path, device="cuda"):
    13→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    14→
    15→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    16→        config = ckpt["config"]
    17→        self.model = LeWMTTS(config).to(self.device)
    18→        self.model.load_state_dict(ckpt["model"])
    19→        self.model.eval()
    20→        self.config = config
    21→
    22→        print("Loading DAC 24kHz...")
    23→        model_path = dac.utils.download(model_type='24khz')
    24→        self.dac_model = dac.DAC.load(model_path)
    25→        self.dac_model.eval()
    26→        self.dac_model = self.dac_model.to(self.device)
    27→
    28→        with torch.no_grad():
    29→            dummy = torch.randn(1, 1, 24000).to(self.device)
    30→            dummy = self.dac_model.preprocess(dummy, 24000)
    31→            z = self.dac_model.encoder(dummy)
    32→            self._start_scale = z.std().item()
    33→        print(f"Model loaded. Start scale: {self._start_scale:.3f}")
    34→
    35→    def text_to_tokens(self, text):
    36→        tokens = list(text.encode("utf-8"))
    37→        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
    38→
    39→    @torch.no_grad()
    40→    def synthesize(self, text, max_steps=500, temperature=0.0):
    41→        text_tokens = self.text_to_tokens(text)
    42→        text_emb = self.model.text_encoder(text_tokens)
    43→
    44→        start_dac = torch.randn(1, 1, 1024, device=self.device) * self._start_scale
    45→        start_emb = self.model.dac_in_proj(start_dac)
    46→
    47→        all_embs = [start_emb]
    48→        prev_norms = []
    49→
    50→        for step in range(max_steps):
    51→            context = torch.cat(all_embs, dim=1)
    52→            next_emb = self.model.predict_next(context, text_emb)
    53→            if temperature > 0:
    54→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
    55→            all_embs.append(next_emb)
    56→
    57→            norm = next_emb.norm().item()
    58→            prev_norms.append(norm)
    59→            if len(prev_norms) > 20:
    60→                recent = prev_norms[-20:]
    61→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
    62→                    break
    63→
    64→        audio_embs = torch.cat(all_embs, dim=1)
    65→        dac_latents = self.model.latents_to_dac(audio_embs)
    66→        dac_latents = dac_latents.transpose(1, 2)  # [1, 1024, T]
    67→
    68→        waveform = self.dac_model.decode(dac_latents)
    69→        return waveform.squeeze().cpu().numpy(), 24000
    70→
    71→    @torch.no_grad()
    72→    def reconstruct(self, audio_path):
    73→        """DAC encode → project through model → DAC decode."""
    74→        wav, sr = torchaudio.load(audio_path)
    75→        wav = wav.unsqueeze(0).to(self.device)
    76→        wav = self.dac_model.preprocess(wav, sr)
    77→
    78→        z = self.dac_model.encoder(wav)  # [1, 1024, T]
    79→        z_in = z.transpose(1, 2)  # [1, T, 1024]
    80→        h = self.model.dac_in_proj(z_in)
    81→        z_out = self.model.latents_to_dac(h)
    82→        z_out = z_out.transpose(1, 2)  # [1, 1024, T]
    83→
    84→        waveform = self.dac_model.decode(z_out)
    85→        return waveform.squeeze().cpu().numpy(), 24000
    86→
    87→    @torch.no_grad()
    88→    def reconstruct_direct(self, audio_path):
    89→        """Pure DAC encode → decode. Baseline quality."""
    90→        wav, sr = torchaudio.load(audio_path)
    91→        wav = wav.unsqueeze(0).to(self.device)
    92→        wav = self.dac_model.preprocess(wav, sr)
    93→        z = self.dac_model.encoder(wav)
    94→        waveform = self.dac_model.decode(z)
    95→        return waveform.squeeze().cpu().numpy(), 24000
    96→
    97→    def save_audio(self, waveform, sr, output_path):
    98→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
    99→        torchaudio.save(output_path, wav_tensor, sr)
   100→        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
   101→
   102→
   103→if __name__ == "__main__":
   104→    parser = argparse.ArgumentParser()
   105→    parser.add_argument("--checkpoint", required=True)
   106→    parser.add_argument("--text", default=None)
   107→    parser.add_argument("--audio", default=None)
   108→    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "dac_baseline"],
   109→                        default="synthesize")
   110→    parser.add_argument("--output", default="output.wav")
   111→    parser.add_argument("--max_steps", type=int, default=500)
   112→    parser.add_argument("--temperature", type=float, default=0.0)
   113→    args = parser.parse_args()
   114→
   115→    tts = LeWMTTSInference(args.checkpoint)
   116→    if args.mode == "synthesize":
   117→        wav, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
   118→    elif args.mode == "reconstruct":
   119→        wav, sr = tts.reconstruct(args.audio)
   120→    elif args.mode == "dac_baseline":
   121→        wav, sr = tts.reconstruct_direct(args.audio)
   122→    tts.save_audio(wav, sr, args.output)
   123→

"""Inference for LeWM TTS v2 (DAC-based)."""

import torch
import torchaudio
import dac
import numpy as np
import argparse
from model import LeWMTTS

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

        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        config = ckpt["config"]
        self.model = LeWMTTS(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()
        self.config = config

        print("Loading DAC 24kHz...")
        model_path = dac.utils.download(model_type='24khz')
        self.dac_model = dac.DAC.load(model_path)
        self.dac_model.eval()
        self.dac_model = self.dac_model.to(self.device)

        with torch.no_grad():
            dummy = torch.randn(1, 1, 24000).to(self.device)
            dummy = self.dac_model.preprocess(dummy, 24000)
            z = self.dac_model.encoder(dummy)
            self._start_scale = z.std().item()
        print(f"Model loaded. Start scale: {self._start_scale:.3f}")

    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=500, temperature=0.0):
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        start_dac = torch.randn(1, 1, 1024, device=self.device) * self._start_scale
        start_emb = self.model.dac_in_proj(start_dac)

        all_embs = [start_emb]
        prev_norms = []

        for step in range(max_steps):
            context = torch.cat(all_embs, dim=1)
            next_emb = self.model.predict_next(context, text_emb)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embs.append(next_emb)

            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_embs, dim=1)
        dac_latents = self.model.latents_to_dac(audio_embs)
        dac_latents = dac_latents.transpose(1, 2)  # [1, 1024, T]

        waveform = self.dac_model.decode(dac_latents)
        return waveform.squeeze().cpu().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """DAC encode → project through model → DAC decode."""
        wav, sr = torchaudio.load(audio_path)
        wav = wav.unsqueeze(0).to(self.device)
        wav = self.dac_model.preprocess(wav, sr)

        z = self.dac_model.encoder(wav)  # [1, 1024, T]
        z_in = z.transpose(1, 2)  # [1, T, 1024]
        h = self.model.dac_in_proj(z_in)
        z_out = self.model.latents_to_dac(h)
        z_out = z_out.transpose(1, 2)  # [1, 1024, T]

        waveform = self.dac_model.decode(z_out)
        return waveform.squeeze().cpu().numpy(), 24000

    @torch.no_grad()
    def reconstruct_direct(self, audio_path):
        """Pure DAC encode → decode. Baseline quality."""
        wav, sr = torchaudio.load(audio_path)
        wav = wav.unsqueeze(0).to(self.device)
        wav = self.dac_model.preprocess(wav, sr)
        z = self.dac_model.encoder(wav)
        waveform = self.dac_model.decode(z)
        return waveform.squeeze().cpu().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)")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None)
    parser.add_argument("--audio", default=None)
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "dac_baseline"],
                        default="synthesize")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--max_steps", type=int, default=500)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSInference(args.checkpoint)
    if args.mode == "synthesize":
        wav, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        wav, sr = tts.reconstruct(args.audio)
    elif args.mode == "dac_baseline":
        wav, sr = tts.reconstruct_direct(args.audio)
    tts.save_audio(wav, sr, args.output)


There are two versions - v1 (mel-based with Griffin-Lim) and v2 (DAC-based). The user said "2 versions" which likely means both v1 and v2 inference scripts. Let me check which checkpoints belong to which version, and also check if v2 has its own checkpoints.


Let me check if v2 has its own checkpoints and find a test text to use.


ls -la /home/ubuntu/lewm-tts/v2/output/ 2>/dev/null || ls -la /home/ubuntu/lewm-tts/v2/ 2>/dev/null

Check v2 output directory


/home/ubuntu/lewm-tts/output/checkpoints/epoch_0070.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0010.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0100.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0020.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0030.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0040.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0150.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0050.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0060.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0070.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0200.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0080.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0090.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0250.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0100.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0110.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0120.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0300.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0130.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0140.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/epoch_0150.pt
/home/ubuntu/lewm-tts/output/checkpoints/epoch_0350.pt
/home/ubuntu/lewm-tts/v2/output/checkpoints/best.pt
/home/ubuntu/lewm-tts/output/checkpoints/best.pt


/home/ubuntu/lewm-tts/v2/processed_data/latents/000000.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000001.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000002.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000003.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000004.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000005.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000006.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000007.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000008.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000009.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000010.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000011.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000012.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000013.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000014.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000015.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000016.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000017.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000018.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000019.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000020.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000021.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000022.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000023.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000024.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000025.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000026.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000027.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000028.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000029.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000030.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000031.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000032.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000033.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000034.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000035.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000036.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000037.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000038.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000039.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000040.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000041.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000042.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000043.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000044.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000045.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000046.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000047.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000048.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000049.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000050.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000051.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000052.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000053.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000054.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000055.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000056.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000057.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000058.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000059.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000060.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000061.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000062.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000063.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000064.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000065.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000066.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000067.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000068.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000069.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000070.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000071.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000072.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000073.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000074.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000075.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000076.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000077.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000078.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000079.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000080.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000081.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000082.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000083.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000084.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000085.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000086.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000087.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000088.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000089.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000090.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000091.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000092.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000093.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000094.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000095.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000096.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000097.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000098.pt
/home/ubuntu/lewm-tts/v2/processed_data/latents/000099.pt
(Results are truncated. Consider using a more specific path or pattern.)


total 1084
drwxrwxr-x 4 ubuntu ubuntu   4096 Mar 23 22:11 .
drwxrwxr-x 5 ubuntu ubuntu   4096 Mar 23 22:05 ..
drwxrwxr-x 2 ubuntu ubuntu   4096 Mar 24 03:38 checkpoints
-rw-rw-r-- 1 ubuntu ubuntu    194 Mar 23 22:11 config.json
drwxrwxr-x 2 ubuntu ubuntu   4096 Mar 23 22:11 logs
-rw-rw-r-- 1 ubuntu ubuntu  91166 Mar 24 03:44 train.log
-rw-rw-r-- 1 ubuntu ubuntu  21182 Mar 23 22:02 v2_ar_synth.wav
-rw-rw-r-- 1 ubuntu ubuntu 480062 Mar 23 22:02 v2_dac_baseline.wav
-rw-rw-r-- 1 ubuntu ubuntu 480062 Mar 23 22:02 v2_reconstruct.wav


So there are two versions:
- v1: latest checkpoints are `epoch_0350.pt` and `best.pt` in `/home/ubuntu/lewm-tts/output/checkpoints/`
- v2: latest checkpoint is `epoch_0150.pt` and `best.pt` in `/home/ubuntu/lewm-tts/v2/output/checkpoints/`

The user said "2 versions latest checkpoints" - so they want inference on:
1. v1 latest: `epoch_0350.pt` (or `best.pt`)
2. v2 latest: `epoch_0150.pt` (or `best.pt`)

I'll run inference on both versions using their latest epoch checkpoints in parallel.


Found both versions:
- **v1** (mel-based): latest checkpoint `epoch_0350.pt`
- **v2** (DAC-based): latest checkpoint `epoch_0150.pt`

Running inference on both with a test sentence:


cd /home/ubuntu/lewm-tts && python inference.py --checkpoint output/checkpoints/epoch_0350.pt --text "Hello, this is a test of the LeWM text to speech system." --output output/v1_epoch350_synth.wav --max_steps 300

Run v1 inference with epoch 350 checkpoint


cd /home/ubuntu/lewm-tts/v2 && python inference.py --checkpoint output/checkpoints/epoch_0150.pt --text "Hello, this is a test of the LeWM text to speech system." --output output/v2_epoch150_synth.wav --max_steps 500

Run v2 inference with epoch 150 checkpoint


Error: Exit code 127
/bin/bash: line 1: python: command not found


which python3 && ls /home/ubuntu/lewm-tts/.venv/bin/python 2>/dev/null || conda info --envs 2>/dev/null | head -10

Find Python executable


cd /home/ubuntu/lewm-tts && python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --text "Hello, this is a test of the LeWM text to speech system." --output output/v1_epoch350_synth.wav --max_steps 300

Run v1 inference with epoch 350 checkpoint


cd /home/ubuntu/lewm-tts/v2 && python3 inference.py --checkpoint output/checkpoints/epoch_0150.pt --text "Hello, this is a test of the LeWM text to speech system." --output output/v2_epoch150_synth.wav --max_steps 500

Run v2 inference with epoch 150 checkpoint


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_epoch350_synth.wav (1.40s)


/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: 139.877
Saved: output/v2_epoch150_synth.wav (2.31s)


Both inference runs completed:

| Version | Checkpoint | Output | Duration |
|---------|-----------|--------|----------|
| **v1** (mel + Griffin-Lim) | `epoch_0350.pt` | `output/v1_epoch350_synth.wav` | 1.40s |
| **v2** (DAC-based) | `epoch_0150.pt` | `v2/output/v2_epoch150_synth.wav` | 2.31s |

Both generated from: *"Hello, this is a test of the LeWM text to speech system."*

You can listen to them to compare quality between v1 and v2.


fuck brother use hindi .. and aso u fucjed up soemwhere ..check meory once on this


The user wants Hindi text for inference and says I messed something up. Let me check my memory file for this project.


     1→# LeWM TTS Project
     2→
     3→## Overview
     4→Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
     5→- Paper: https://le-wm.github.io/
     6→- Project folder: `/home/ubuntu/lewm-tts`
     7→
     8→## Architecture (Implemented)
     9→- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
    10→- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
    11→- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text
    12→- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents
    13→- **Vocoder**: Griffin-Lim (fallback; HiFi-GAN planned)
    14→- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)
    15→- **Model size**: 14.50M parameters
    16→- **KL annealing**: ramps up over first 20% of training
    17→
    18→## Files
    19→- `preprocess.py` — segments audio, extracts mels, saves manifest
    20→- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, LeWMTTS
    21→- `dataset.py` — TTSDataset, collate_fn, build_dataloader
    22→- `train.py` — training loop with logging, checkpointing, KL annealing
    23→- `inference.py` — AR generation + Griffin-Lim vocoder
    24→
    25→## Dataset
    26→- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
    27→- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)
    28→- Mel config: n_fft=1024, hop=256, n_mels=80, sr=24000
    29→- Segment stats: min=3s, max=14.4s, mean=9.2s, 99.3% fit in 1200 mel frames
    30→
    31→## Training Results (3 epoch test)
    32→- Prediction loss: 1.54 → 0.017 (99% drop — excellent)
    33→- Recon loss: 2.25 → 0.68 (good convergence)
    34→- KL loss: stabilizes ~1.15 with kl_weight=0.1 + annealing
    35→- Speed: ~27s/epoch on A100 80GB, batch_size=16
    36→- No NaNs or instabilities
    37→
    38→## Recommended Training Command
    39→```
    40→cd /home/ubuntu/lewm-tts
    41→python3 train.py --epochs 200 --batch_size 16 --lr 1e-4 --kl_weight 0.1
    42→```
    43→
    44→## Status
    45→- [x] Project folder created
    46→- [x] Architecture implemented & tested
    47→- [x] Data preprocessing pipeline (7,009 segments)
    48→- [x] Model implementation (14.50M params)
    49→- [x] Training loop verified (stable, losses converge)
    50→- [x] Inference pipeline (AR + Griffin-Lim)
    51→- [ ] Full training run (200 epochs, ~1.5h on A100)
    52→- [ ] Quality evaluation & audio samples
    53→- [ ] HiFi-GAN vocoder integration (for better audio quality)
    54→

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`

## Architecture (Implemented)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text
- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents
- **Vocoder**: Griffin-Lim (fallback; HiFi-GAN planned)
- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)
- **Model size**: 14.50M parameters
- **KL annealing**: ramps up over first 20% of training

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing
- `inference.py` — AR generation + Griffin-Lim vocoder

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=80, sr=24000
- Segment stats: min=3s, max=14.4s, mean=9.2s, 99.3% fit in 1200 mel frames

## Training Results (3 epoch test)
- Prediction loss: 1.54 → 0.017 (99% drop — excellent)
- Recon loss: 2.25 → 0.68 (good convergence)
- KL loss: stabilizes ~1.15 with kl_weight=0.1 + annealing
- Speed: ~27s/epoch on A100 80GB, batch_size=16
- No NaNs or instabilities

## Recommended Training Command
```
cd /home/ubuntu/lewm-tts
python3 train.py --epochs 200 --batch_size 16 --lr 1e-4 --kl_weight 0.1
```

## Status
- [x] Project folder created
- [x] Architecture implemented & tested
- [x] Data preprocessing pipeline (7,009 segments)
- [x] Model implementation (14.50M params)
- [x] Training loop verified (stable, losses converge)
- [x] Inference pipeline (AR + Griffin-Lim)
- [ ] Full training run (200 epochs, ~1.5h on A100)
- [ ] Quality evaluation & audio samples
- [ ] HiFi-GAN vocoder integration (for better audio quality)


Right, the dataset is **Hindi** (Modi speeches). My bad — I used English text. Let me rerun with Hindi text. Also the dataset is Hindi so the model was trained on Hindi audio.


cd /home/ubuntu/lewm-tts && python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --text "नमस्ते, यह लेडब्ल्यूएम टेक्स्ट टू स्पीच सिस्टम का परीक्षण है।" --output output/v1_epoch350_hindi.wav --max_steps 300

Run v1 inference with Hindi text


python3 inference.py --checkpoint output/checkpoints/epoch_0150.pt --text "नमस्ते, यह लेडब्ल्यूएम टेक्स्ट टू स्पीच सिस्टम का परीक्षण है।" --output output/v2_epoch150_hindi.wav --max_steps 500

Run v2 inference with Hindi text


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_epoch350_hindi.wav (2.25s)


Model loaded from output/checkpoints/epoch_0150.pt
Saved: output/v2_epoch150_hindi.wav (4.21s)


Wait — v2 inference ran from the wrong directory. Let me fix that.


cd /home/ubuntu/lewm-tts/v2 && python3 inference.py --checkpoint output/checkpoints/epoch_0150.pt --text "नमस्ते, यह लेडब्ल्यूएम टेक्स्ट टू स्पीच सिस्टम का परीक्षण है।" --output output/v2_epoch150_hindi.wav --max_steps 500

Run v2 inference from correct v2 directory


/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: 134.814
Saved: output/v2_epoch150_hindi.wav (2.12s)


Done. Reran with Hindi text:

| Version | Checkpoint | Output | Duration |
|---------|-----------|--------|----------|
| **v1** (mel + Griffin-Lim) | `epoch_0350` | `lewm-tts/output/v1_epoch350_hindi.wav` | 2.25s |
| **v2** (DAC-based) | `epoch_0150` | `lewm-tts/v2/output/v2_epoch150_hindi.wav` | 2.12s |

Text used: *"नमस्ते, यह लेडब्ल्यूएम टेक्स्ट टू स्पीच सिस्टम का परीक्षण है।"*

Sorry about the English text earlier — the model was trained on Hindi (Modi dataset) so it should only be tested with Hindi input.


brother we got audoio at 25th epoch for v1 .. y it not cmng now ..u fucking up somehere ...check fuckingmmeoery


The user is saying we got audio at epoch 25 for v1 but now it's not working properly. Let me check the memory more carefully and look at what happened previously.


Let me dig deeper into what worked before.


ls -la /home/ubuntu/lewm-tts/output/*.wav 2>/dev/null

List existing wav files in v1 output


-rw-rw-r-- 1 ubuntu ubuntu 411214 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_ar_synth.wav
-rw-rw-r-- 1 ubuntu ubuntu 456782 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_groundtruth_gl.wav
-rw-rw-r-- 1 ubuntu ubuntu 456782 Mar 23 21:16 /home/ubuntu/lewm-tts/output/test_groundtruth_vocos.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_predict_tf.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:11 /home/ubuntu/lewm-tts/output/test_real_griffinlim.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:11 /home/ubuntu/lewm-tts/output/test_recon_griffinlim.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_reconstruct.wav
-rw-rw-r-- 1 ubuntu ubuntu 613966 Mar 23 21:09 /home/ubuntu/lewm-tts/output/test_sample_epoch10.wav
-rw-rw-r-- 1 ubuntu ubuntu 479822 Mar 23 21:46 /home/ubuntu/lewm-tts/output/v1_best_groundtruth.wav
-rw-rw-r-- 1 ubuntu ubuntu 478798 Mar 23 21:46 /home/ubuntu/lewm-tts/output/v1_best_predict_tf.wav
-rw-rw-r-- 1 ubuntu ubuntu 478798 Mar 23 21:46 /home/ubuntu/lewm-tts/output/v1_best_reconstruct.wav
-rw-rw-r-- 1 ubuntu ubuntu 108110 Mar 24 03:47 /home/ubuntu/lewm-tts/output/v1_epoch350_hindi.wav
-rw-rw-r-- 1 ubuntu ubuntu  67150 Mar 24 03:45 /home/ubuntu/lewm-tts/output/v1_epoch350_synth.wav
-rw-rw-r-- 1 ubuntu ubuntu 202318 Mar 24 03:47 /home/ubuntu/lewm-tts/output/v2_epoch150_hindi.wav
-rw-rw-r-- 1 ubuntu ubuntu 456782 Mar 23 21:16 /home/ubuntu/lewm-tts/output/vocos_groundtruth.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:16 /home/ubuntu/lewm-tts/output/vocos_predict_tf.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:16 /home/ubuntu/lewm-tts/output/vocos_reconstruct.wav


/home/ubuntu/.claude/projects/-home-ubuntu/memory/MEMORY.md
/home/ubuntu/.claude/projects/-home-ubuntu/memory/lewm-tts.md


I see — previously we generated audio using **reconstruct** and **predict** modes (teacher-forced with real mel input), not pure synthesis. Those worked well because they use real audio as input. Let me check the old working commands and the previous inference approach.


ls -la /home/ubuntu/lewm-tts/output/test_sample_epoch10.wav /home/ubuntu/lewm-tts/output/test_reconstruct.wav /home/ubuntu/lewm-tts/output/test_ar_synth.wav

Check sizes of old working samples


ls /home/ubuntu/lewm-tts/processed_data/mels/ 2>/dev/null | head -5

Find mel files for reconstruction


-rw-rw-r-- 1 ubuntu ubuntu 411214 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_ar_synth.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 23 21:12 /home/ubuntu/lewm-tts/output/test_reconstruct.wav
-rw-rw-r-- 1 ubuntu ubuntu 613966 Mar 23 21:09 /home/ubuntu/lewm-tts/output/test_sample_epoch10.wav


000000.pt
000001.pt
000002.pt
000003.pt
000004.pt


Now I see. The old working audio files were **much larger** (411-613KB) vs what I just generated (67-108KB). The previous approach used **reconstruction** mode with real mel spectrograms as input, and the AR synth also produced decent output at earlier epochs.

Let me run all 3 modes like before — reconstruct, predict (teacher-forced), and synthesize — using the same mel file:


cd /home/ubuntu/lewm-tts && python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --mel processed_data/mels/000000.pt --mode reconstruct --output output/v1_e350_reconstruct.wav

v1 reconstruct mode with real mel


cd /home/ubuntu/lewm-tts && python3 inference.py --checkpoint output/checkpoints/epoch_0350.pt --mel processed_data/mels/000000.pt --mode predict --output output/v1_e350_predict.wav

v1 predict (teacher-forced) mode


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

v1 synthesize mode with Hindi


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_e350_reconstruct.wav (9.50s)


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_e350_predict.wav (9.50s)


Model loaded from output/checkpoints/epoch_0350.pt
Saved: output/v1_e350_synth.wav (3.57s)


Now v2:


ls /home/ubuntu/lewm-tts/v2/processed_data/audio/ 2>/dev/null | head -3 || ls /home/ubuntu/lewm-tts/processed_data/audio/ 2>/dev/null | head -3

Find audio files for v2 reconstruction


(Bash completed with no output)


ls /home/ubuntu/lewm-tts/v2/processed_data/ && find /home/ubuntu/lewm-tts/processed_data -name "*.wav" | head -3

Find audio/wav files for v2


config.json
latents
manifest.json


find /home/ubuntu/modi_dataset -name "*.wav" | head -3

Find raw audio files


find /home/ubuntu/modi_dataset -type f | head -5 && find /home/ubuntu/lewm-tts/processed_data -name "*.wav" -o -name "*.flac" -o -name "*.mp3" | head -3

Find any audio files in dataset


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 with Hindi text


/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: 139.260
Saved: output/v2_e150_synth.wav (6.68s)


Now we have proper outputs. Here's the full comparison:

**v1 (mel + Griffin-Lim) — epoch 350:**
| Mode | File | Duration |
|------|------|----------|
| Reconstruct | `output/v1_e350_reconstruct.wav` | 9.50s |
| Predict (teacher-forced) | `output/v1_e350_predict.wav` | 9.50s |
| Synthesize (Hindi) | `output/v1_e350_synth.wav` | 3.57s |

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

The reconstruct/predict modes produce proper length audio (9.5s) matching the training segments. The pure AR synthesis is shorter — that's the AR stopping early. V2 synthesis gave 6.68s which is better than the 2s I got before with English text.


──────────────────────────┼────────────────────────────────┼──────────┤
  │ Synthesize (Hindi)       │ output/v1_e350_synth.wav       │  not good ... we ddidnt used giffin-lim we used vococs


The user is saying we used Vocos, not Griffin-Lim. Let me check the memory and the existing wav files - there were vocos_*.wav files in the output. Let me look at how inference was done before with Vocos.


Right, I see `vocos_*.wav` files in the output from before. Let me check how Vocos was used.


grep -r "vocos" /home/ubuntu/lewm-tts/*.py /home/ubuntu/lewm-tts/v2/*.py 2>/dev/null

Find Vocos usage in code


/home/ubuntu/lewm-tts/output/vocos_groundtruth.wav
/home/ubuntu/lewm-tts/output/vocos_reconstruct.wav
/home/ubuntu/lewm-tts/output/vocos_predict_tf.wav


Check if Vocos is installed


vocos                                    0.1.0

vocos                                    0.1.0


Vocos is installed. The inference.py currently uses Griffin-Lim but previously we used Vocos as the vocoder. Let me update the inference script to use Vocos instead.


python3 -c "from vocos import Vocos; v = Vocos.from_pretrained('charactr/vocos-mel-24khz'); print('Vocos loaded OK')"

Test Vocos import and loading


Vocos loaded OK


Good. Now let me write a quick inference using Vocos instead of Griffin-Lim:


    56→        self.model.load_state_dict(ckpt["model"])
    57→        self.model.eval()
    58→
    59→        # Vocoder
    60→        self.vocoder = GriffinLimVocoder(
    61→            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
    62→        )
    63→
    64→        self.config = config
    65→

        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # Vocoder
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

        self.config = config


ToolSearch

Question: select:Edit


Question: select:Edit


Edit

        # Vocoder
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

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


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

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

"""
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
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

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

-        self.vocoder = GriffinLimVocoder(

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

+        self.vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz")


Now update the vocoder call — Vocos takes mel differently than Griffin-Lim. Let me check Vocos API:


Bash

python3 -c "from vocos import Vocos; help(Vocos.decode)" 2>&1 | head -20

Check Vocos decode API


Help on function decode in module vocos.pretrained:

decode(self, features_input: 'torch.Tensor', **kwargs: 'Any') -> 'torch.Tensor'
    Method to decode audio waveform from already calculated features. The features input is passed through
    the backbone and the head to reconstruct the audio output.

    Args:
        features_input (Tensor): The input tensor of features of shape (B, C, L), where B is the batch size,
                                 C denotes the feature dimension, and L is the sequence length.

    Returns:
        Tensor: The output tensor representing the reconstructed audio waveform of shape (B, T).


Vocos expects `(B, C, L)` — that's `(batch, n_mels, time)`. Let me check how the mel decoder outputs and how the vocoder is called:


   110→
   111→        # Decode to mel
   112→        mel = self.model.mel_decoder(audio_embs)
   113→
   114→        # Vocoder
   115→        waveform = self.vocoder(mel.cpu())
   116→        return waveform.squeeze().numpy(), 24000
   117→
   118→    @torch.no_grad()
   119→    def reconstruct(self, mel_path):
   120→        """
   121→        Teacher-forced reconstruction: encode real mel → decode back.
   122→        Tests mel decoder + vocoder quality without AR drift.
   123→        """
   124→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
   125→        # [1, n_mels, T]
   126→
   127→        # Encode
   128→        z, mu, logvar = self.model.audio_encoder(mel)
   129→
   130→        # Decode from mean embeddings
   131→        mel_recon = self.model.mel_decoder(mu)
   132→
   133→        # Vocoder
   134→        waveform = self.vocoder(mel_recon.cpu())
   135→        return waveform.squeeze().numpy(), 24000
   136→
   137→    @torch.no_grad()
   138→    def reconstruct_with_prediction(self, mel_path):
   139→        """

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

        # Vocoder
        waveform = self.vocoder(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(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """


   150→        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
   151→        text_emb = self.model.text_encoder(text_tokens)
   152→
   153→        # Run predictor on real embeddings (teacher-forced)
   154→        predicted = self.model.predictor(mu[:, :-1], text_emb)
   155→
   156→        # Use real first embedding + predicted rest
   157→        combined = torch.cat([mu[:, :1], predicted], dim=1)
   158→
   159→        # Decode
   160→        mel_recon = self.model.mel_decoder(combined)
   161→        waveform = self.vocoder(mel_recon.cpu())
   162→        return waveform.squeeze().numpy(), 24000
   163→
   164→    def save_audio(self, waveform, sr, output_path):
   165→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
   166→        torchaudio.save(output_path, wav_tensor, sr)
   167→        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
   168→
   169→

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


I need to check what shape mel_decoder outputs. Let me check quickly:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model import LeWMTTS
ckpt = torch.load('output/checkpoints/epoch_0350.pt', map_location='cpu', weights_only=False)
m = LeWMTTS(ckpt['config'])
m.load_state_dict(ckpt['model'])
dummy = torch.randn(1, 10, ckpt['config']['d_model'])
out = m.mel_decoder(dummy)
print('mel_decoder output shape:', out.shape)
"

Check mel_decoder output shape


[W324 03:53:09.016347435 NNPACK.cpp:56] Could not initialize NNPACK! Reason: Unsupported hardware.
mel_decoder output shape: torch.Size([1, 80, 40])


Output is `(B, 80, T)` — exactly what Vocos expects `(B, C, L)`. Now I need to replace the `self.vocoder(mel.cpu())` calls with `self.vocoder.decode(mel.cpu())`:


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.vocoder = 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(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(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(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(mel.cpu())

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


"""
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.vocoder = 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(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(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(mel_recon.cpu())

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


Now rerun inference with Vocos:


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 vocoder


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 vocoder


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


Error: Exit code 1
Model loaded from output/checkpoints/epoch_0350.pt
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/inference.py", line 198, in <module>
    main()
  File "/home/ubuntu/lewm-tts/inference.py", line 189, in main
    waveform, sr = tts.reconstruct(args.mel)
  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 134, in reconstruct
    waveform = self.vocoder.decode(mel_recon.cpu())
  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/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 112, in decode
    x = self.backbone(features_input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/models.py", line 79, in forward
    x = self.embed(x)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 371, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 366, in _conv_forward
    return F.conv1d(
RuntimeError: Given groups=1, weight of size [512, 100, 7], expected input[1, 80, 892] to have 100 channels, but got 80 channels instead


Error: Exit code 1
Model loaded from output/checkpoints/epoch_0350.pt
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/inference.py", line 198, in <module>
    main()
  File "/home/ubuntu/lewm-tts/inference.py", line 192, in main
    waveform, sr = tts.reconstruct_with_prediction(args.mel)
  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 161, in reconstruct_with_prediction
    waveform = self.vocoder.decode(mel_recon.cpu())
  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/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 112, in decode
    x = self.backbone(features_input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/models.py", line 79, in forward
    x = self.embed(x)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 371, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 366, in _conv_forward
    return F.conv1d(
RuntimeError: Given groups=1, weight of size [512, 100, 7], expected input[1, 80, 892] to have 100 channels, but got 80 channels instead


Error: Exit code 1
Model loaded from output/checkpoints/epoch_0350.pt
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/inference.py", line 198, in <module>
    main()
  File "/home/ubuntu/lewm-tts/inference.py", line 186, 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 115, in synthesize
    waveform = self.vocoder.decode(mel.cpu())
  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/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 112, in decode
    x = self.backbone(features_input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/models.py", line 79, in forward
    x = self.embed(x)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 371, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 366, in _conv_forward
    return F.conv1d(
RuntimeError: Given groups=1, weight of size [512, 100, 7], expected input[1, 80, 260] to have 100 channels, but got 80 channels instead


Vocos mel-24khz expects 100 channels, not 80. The model was trained with 80 mels. The right Vocos model must be different — let me check:


python3 -c "
from vocos import Vocos
# Try the correct model for 80 mels
v = Vocos.from_pretrained('BSC-LT/vocos-mel-22khz')
print('22khz model loaded')
" 2>&1; python3 -c "
from vocos import Vocos
import inspect
# Check what models are available
src = inspect.getfile(Vocos)
print('Vocos source:', src)
" 2>&1

Find correct Vocos model for 80 mels


Traceback (most recent call last):
  File "<string>", line 4, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 69, in from_pretrained
    model = cls.from_hparams(config_path)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 56, in from_hparams
    feature_extractor = instantiate_class(args=(), init=config["feature_extractor"])
  File "/home/ubuntu/.local/lib/python3.10/site-packages/vocos/pretrained.py", line 30, in instantiate_class
    return args_class(*args, **kwargs)
TypeError: MelSpectrogramFeatures.__init__() got an unexpected keyword argument 'f_min'
Vocos source: /home/ubuntu/.local/lib/python3.10/site-packages/vocos/pretrained.py


python3 -c "
from vocos import Vocos
v = Vocos.from_pretrained('charactr/vocos-mel-24khz')
print('n_mels expected:', v.feature_extractor.n_mels if hasattr(v, 'feature_extractor') else 'unknown')
print('backbone embed weight shape:', v.backbone.embed.weight.shape)
"

Check Vocos 24khz model details


Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 4, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1964, in __getattr__
    raise AttributeError(
AttributeError: 'MelSpectrogramFeatures' object has no attribute 'n_mels'


python3 -c "
from vocos import Vocos
v = Vocos.from_pretrained('charactr/vocos-mel-24khz')
print(v.backbone.embed)
print(v.feature_extractor)
"

Check Vocos architecture details


Conv1d(100, 512, kernel_size=(7,), stride=(1,), padding=(3,))
MelSpectrogramFeatures(
  (mel_spec): MelSpectrogram(
    (spectrogram): Spectrogram()
    (mel_scale): MelScale()
  )
)


Vocos mel-24khz expects 100 mels. Our model outputs 80. We need to either pad to 100 or use a linear projection. The simplest fix — pad the 80 mels to 100:


Edit

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

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


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

"""
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.vocoder = 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()

-        self.vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz")

+        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")


Now add a vocoder wrapper method and replace calls:


        self.config = config

        # Compute start token from training data statistics

        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