        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return z, mu, logvar

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

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

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

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

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

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

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

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

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

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

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

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

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

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

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

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

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

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

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

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

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

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

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

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

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

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

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

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

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

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

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

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

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

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

        # ─── Loss 4: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

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

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

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

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

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

+        downsample_factor = config.get("downsample_factor", 4)

-            num_layers=audio_layers, dropout=dropout,

+            num_layers=audio_layers, downsample_factor=downsample_factor,


Edit

self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)

self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return z, mu, logvar

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

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

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

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

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

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

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

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

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

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

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

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

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

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

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

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

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

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

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

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

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

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

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

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

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

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

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

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

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

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

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

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

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

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

        # ─── Loss 4: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

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

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

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

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

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

-        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)

+        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)


Now add `downsample_factor` to the training args and config:


        "spectral_weight": args.spectral_weight,
    }

        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
    }


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

        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
    }

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            optimizer.step()
            scheduler.step()

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

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

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

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

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

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

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

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

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

+        "downsample_factor": args.downsample_factor,


    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)


    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--downsample_factor", type=int, default=4)

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

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "downsample_factor": args.downsample_factor,
    }
    model, config = build_model(config)
    model = model.to(device)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            optimizer.step()
            scheduler.step()

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

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

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

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

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

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

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

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

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

+    parser.add_argument("--downsample_factor", type=int, default=4)


Let me verify it works:


Bash

python3 -c "
from model import LeWMTTS, count_parameters

# Test 2x downsample
config_2x = {'d_model': 256, 'nhead': 4, 'n_mels': 100, 'text_vocab_size': 256,
             'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
             'dropout': 0.1, 'downsample_factor': 2}
model_2x = LeWMTTS(config_2x)
print(f'2x downsample: {count_parameters(model_2x)/1e6:.2f}M')

# Test 4x still works (backward compat)
config_4x = {'d_model': 256, 'nhead': 4, 'n_mels': 100, 'text_vocab_size': 256,
             'text_encoder_layers': 4, 'audio_encoder_layers': 4, 'predictor_layers': 6,
             'dropout': 0.1}
model_4x = LeWMTTS(config_4x)
print(f'4x downsample: {count_parameters(model_4x)/1e6:.2f}M')

# Test forward pass with 2x
import torch
mel = torch.randn(2, 100, 400)
text = torch.randint(0, 256, (2, 20))
out = model_2x(mel, text)
print(f'2x forward OK: recon_loss={out[\"recon_loss\"].item():.4f}')
" 2>&1

Verify 2x downsample model works


2x downsample: 21.11M
4x downsample: 24.00M
2x forward OK: recon_loss=0.7944


Ready. **21.1M params** with 2x downsample, backward compatible with 4x.

v3b is still running (won't be affected — already loaded in memory). Want me to start the 2x downsample training now in parallel as `output_v3c`?


yeah


mkdir -p /home/ubuntu/lewm-tts/output_v3c && python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3c --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --downsample_factor 2 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10 > /home/ubuntu/lewm-tts/train_v3c.log 2>&1

Start v3c training with 2x downsample


Command running in background with ID: b74eiv52u. Output is being written to: /tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/b74eiv52u.output


sleep 35 && head -8 /home/ubuntu/lewm-tts/train_v3c.log

Check v3c first epoch


<task-notification>
<task-id>b74eiv52u</task-id>
<tool-use-id>toolu_vrtx_01HQmq5PuUo1ZhKFaYaCaZUC</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/b74eiv52u.output</output-file>
<status>failed</status>
<summary>Background command "Start v3c training with 2x downsample" failed with exit code 1</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/b74eiv52u.output


Device: cuda
LeWM TTS model: 21.11M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/train.py", line 229, in <module>
    train(args)
  File "/home/ubuntu/lewm-tts/train.py", line 109, in train


Check full error


Device: cuda
LeWM TTS model: 21.11M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Traceback (most recent call last):
  File "/home/ubuntu/lewm-tts/train.py", line 229, in <module>
    train(args)
  File "/home/ubuntu/lewm-tts/train.py", line 109, in train
    losses = model(mel, text_tokens, mel_mask, text_mask)
  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/lewm-tts/model.py", line 537, in forward
    z, mu, logvar = self.audio_encoder(mel, mel_mask)
  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/lewm-tts/model.py", line 126, in forward
    x = self.transformer(x, src_key_padding_mask=mel_mask)
  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/transformer.py", line 524, in forward
    output = mod(
  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/transformer.py", line 935, in forward
    + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py", line 949, in _sa_block
    x = self.self_attn(
  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/activation.py", line 1488, in forward
    attn_output, attn_output_weights = F.multi_head_attention_forward(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/functional.py", line 6307, in multi_head_attention_forward
    q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/functional.py", line 5706, in _in_projection_packed
    .contiguous()
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 446.00 MiB. GPU 0 has a total capacity of 79.25 GiB of which 65.19 MiB is free. Process 3530968 has 65.61 GiB memory in use. Including non-PyTorch memory, this process has 13.57 GiB memory in use. Of the allocated memory 12.90 GiB is allocated by PyTorch, and 184.39 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


OOM — v3b is still using the GPU. Two models can't fit at batch_size=256. Let me check how much v3b has left and reduce batch size for v3c.


Check v3b progress


  [78/200] step 2100 | loss=2.3289 pred=0.0195 kl=2.4343 recon=0.6279 spec=0.6862 | grad=5.88 | 21.0s
Epoch 78/200 | avg_loss=2.3364 pred=0.0193 kl=2.4395 recon=0.6274 spec=0.6919 | 27.0s


v3b is at epoch 78/200 — still 55 min left. The 2x downsample doubles sequence length so both can't run together at batch 256.

Options:
1. **Kill v3b** and start v3c full speed
2. **Wait** for v3b to finish, then start v3c

Given recon is only at 0.63 and barely improving, I'd say kill v3b — we already know 4x downsample sounds robotic. Your call?


Kill v3b and start v3c?


first run a sample on v3b


python3 -c "
import torch, glob, torchaudio, json
from model import LeWMTTS
from vocos import Vocos

ckpt = torch.load('output_v3b/checkpoints/best.pt', map_location='cuda', weights_only=False)
model = LeWMTTS(ckpt['config']).cuda()
model.load_state_dict(ckpt['model'])
model.eval()
print(f'v3b epoch {ckpt[\"epoch\"]}')

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

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))
with open('processed_data_100mel/manifest.json') as f:
    manifest = json.load(f)

with torch.no_grad():
    for idx in [0, 100, 500]:
        mel = torch.load(mel_files[idx], weights_only=True).unsqueeze(0).cuda()
        z, mu, logvar = model.audio_encoder(mel)
        mel_recon = model.mel_decoder(mu)
        wav = vocos.decode(mel_recon)
        torchaudio.save(f'sample_recon_v3b_{idx}.wav', wav.cpu(), 24000)
        T = min(mel.shape[2], mel_recon.shape[2])
        l1 = (mel[:,:,:T] - mel_recon[:,:,:T]).abs().mean().item()
        print(f'Recon {idx}: L1={l1:.4f}')

    # Teacher-forced
    for idx in [0, 100]:
        mel = torch.load(mel_files[idx], weights_only=True).unsqueeze(0).cuda()
        text = manifest[idx]['text']
        text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()
        z, mu, logvar = model.audio_encoder(mel)
        text_emb = model.text_encoder(text_tokens)
        predicted = model.predictor(mu[:, :-1], text_emb)
        combined = torch.cat([mu[:, :1], predicted], dim=1)
        mel_pred = model.mel_decoder(combined)
        wav = vocos.decode(mel_pred)
        torchaudio.save(f'sample_tf_v3b_{idx}.wav', wav.cpu(), 24000)
        print(f'TF {idx} saved')

    # AR with voice prompt
    mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0).cuda()
    text = manifest[0]['text']
    text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()
    text_emb = model.text_encoder(text_tokens)
    z, mu, logvar = model.audio_encoder(mel)
    prompt_len = 20
    all_embs = [mu[:, :prompt_len]]
    cache = model.init_ar_cache()
    for i in range(prompt_len):
        _, cache = model.predict_next_cached(mu[:, i:i+1], text_emb, i, cache)
    emb = mu[:, prompt_len-1:prompt_len]
    for step in range(prompt_len, 300):
        emb, cache = model.predict_next_cached(emb, text_emb, step, cache)
        all_embs.append(emb)
    mel_ar = model.mel_decoder(torch.cat(all_embs, dim=1))
    wav_ar = vocos.decode(mel_ar)
    torchaudio.save('sample_ar_prompted_v3b.wav', wav_ar.cpu(), 24000)
    print('AR prompted saved')
print('Done!')
" 2>&1

Generate all v3b samples at epoch 78


v3b epoch 90
Recon 0: L1=0.5639
Recon 100: L1=0.5708
Recon 500: L1=0.5813
TF 0 saved
TF 100 saved
AR prompted saved
Done!


Epoch 90 samples ready. L1 dropped nicely to **0.56-0.58** (was 0.64 at epoch 57).

| File | Description |
|------|-------------|
| `sample_recon_v3b_0/100/500.wav` | Reconstruction (L1≈0.57) |
| `sample_tf_v3b_0/100.wav` | Teacher-forced prediction |
| `sample_ar_prompted_v3b.wav` | AR with voice prompt |

How do they sound now?


stop and run it .. but i think we fucndamenatlly dng soemthing wrong in traing or infernce


The user wants to stop v3b and start v3c, but they also think there's something fundamentally wrong. Let me think about what could be fundamentally wrong.

Let me reconsider the architecture and training:

1. The model trains with: mel → AudioEncoder → z → Predictor predicts z[t+1] from z[t] + text → MelDecoder reconstructs mel from z
2. During training, the MelDecoder receives z (sampled from the encoder), and its loss is to reconstruct the input mel
3. During inference reconstruction, we use mu (not sampled z)

Wait... actually there might be a fundamental issue. Let me re-read the forward method.

Looking at the code, during training the decoder uses the sampled z, but at inference time it uses mu instead. However, at eval time the AudioEncoder returns z=mu with no sampling, so that should be consistent. Let me think about what other fundamental issues could be at play here.

I'm wondering if there's a mismatch in how we're using the Vocos vocoder—it expects log-mel spectrograms in a specific format, and there might be differences in how we're preprocessing them. The `safe_log` function in Vocos could be different from our clamping approach, and there's also the question of whether Vocos uses different padding settings like `center=True` that we're not accounting for.

Actually, I realize the key issue: Vocos's `decode()` method expects features that match what its backbone was trained on—which is log-mel spectrograms from its own feature extractor pipeline. Since we matched the mel config (sr=24000, n_fft=1024, hop=256, n_mels=100, power=1), the only real differences would be the `center` padding and the exact log clamp value, which shouldn't cause fundamental issues. The user tested ground truth mels through the decoder...

The real problem is likely that the model's error is structured rather than random—spectral blur and missing formants from the 4x downsampling bottleneck and L1/MSE loss pushing toward mean predictions. Even though random noise at L1=0.84 sounded natural in testing, the model's learned error pattern at L1=0.56 is fundamentally different. I should check if there's a normalization mismatch in how Vocos processes the features. Both Vocos and our preprocessing use center=True for the mel spectrograms, so that's not the issue. The real problem might just be that 256 dimensions per latent step isn't sufficient capacity, or the architecture itself needs a fundamental redesign. I'll stop v3b and move forward with v3c as requested while flagging this concern.


ps aux | grep "train.py" | grep -v grep | awk '{print $2}'

Find v3b training PID


3530948
3530968
3554029
3554030
3554031
3554032


Kill v3b training


<task-notification>
<task-id>bx8o2n8dg</task-id>
<tool-use-id>toolu_vrtx_01WTEeL8KYhZgpW9D7GZV9Qu</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/bx8o2n8dg.output</output-file>
<status>failed</status>
<summary>Background command "Start v3b training with kl=0.01, recon=2.0, spectral=1.5" failed with exit code 144</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/bx8o2n8dg.output


(Bash completed with no output)


sleep 3 && nvidia-smi --query-gpu=memory.free --format=csv,noheader

Check GPU memory freed


13965 MiB


sleep 5 && nvidia-smi --query-gpu=memory.free --format=csv,noheader

Wait for full GPU memory release


Check if training processes are gone


ubuntu   3530968  102  2.2 85879172 2805508 ?    Rl   06:20  43:42 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3b --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3554329 15.0  1.0 85915584 1263480 ?    Sl   07:02   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3b --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3554330 15.0  1.0 85914804 1262472 ?    Sl   07:02   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3b --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3554331 12.3  1.0 85914660 1262516 ?    Sl   07:02   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3b --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3554332 12.1  1.0 85912980 1260752 ?    Sl   07:02   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3b --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10


Force kill all v3b processes


sleep 3 && nvidia-smi --query-gpu=memory.free --format=csv,noheader

Confirm GPU freed


81156 MiB


GPU free. Now before blindly starting v3c, you're right — let me check if something fundamental is wrong. Let me verify the **GT mel → Vocos → audio** path vs **GT mel → encoder → decoder → Vocos** to see exactly where quality breaks:


python3 -c "
import torch, glob, torchaudio
from vocos import Vocos

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

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))
mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0).cuda()

# 1. What does Vocos produce from its OWN mel extraction?
audio_file = 'processed_data_100mel/audio/000000.wav'
wav_orig, sr = torchaudio.load(audio_file)
print(f'Original audio: sr={sr}, len={wav_orig.shape[1]/sr:.2f}s')

# Vocos end-to-end (audio → vocos mel → vocos decode)
wav_orig_gpu = wav_orig.cuda()
with torch.no_grad():
    vocos_mel = vocos.feature_extractor(wav_orig_gpu)
    wav_vocos_roundtrip = vocos.decode(vocos_mel)
    torchaudio.save('test_vocos_roundtrip.wav', wav_vocos_roundtrip.cpu(), 24000)
    print(f'Vocos roundtrip mel shape: {vocos_mel.shape}, range: [{vocos_mel.min():.2f}, {vocos_mel.max():.2f}]')

# 2. Our preprocessed mel → Vocos decode
with torch.no_grad():
    wav_our_mel = vocos.decode(mel)
    torchaudio.save('test_our_mel_vocos.wav', wav_our_mel.cpu(), 24000)
    print(f'Our mel shape: {mel.shape}, range: [{mel.min():.2f}, {mel.max():.2f}]')

# 3. Compare the two mels
T = min(vocos_mel.shape[2], mel.shape[2])
mel_diff = (vocos_mel[:,:,:T] - mel[:,:,:T]).abs().mean().item()
print(f'Mel diff (vocos vs ours): {mel_diff:.4f}')

# 4. Check safe_log vs our log
import inspect
# Find safe_log
from vocos.feature_extractors import safe_log
print(f'Vocos safe_log: {inspect.getsource(safe_log)}')
" 2>&1

Compare Vocos mel extraction vs our preprocessing


Original audio: sr=24000, len=9.52s
Vocos roundtrip mel shape: torch.Size([1, 100, 893]), range: [-8.69, 4.46]
Our mel shape: torch.Size([1, 100, 893]), range: [-8.69, 4.46]
Mel diff (vocos vs ours): 0.0000
Vocos safe_log: def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
    """
    Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.

    Args:
        x (Tensor): Input tensor.
        clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.

    Returns:
        Tensor: Element-wise logarithm of the input tensor with clipping applied.
    """
    return torch.log(torch.clip(x, min=clip_val))


Good news — **mel diff is 0.0000**. Our preprocessing matches Vocos exactly. So the mel pipeline is not the problem.

The issue is purely that the mel decoder at L1=0.56 still produces structured errors (blurred formants). Listen to `test_vocos_roundtrip.wav` and `test_our_mel_vocos.wav` — they should both sound perfect (confirming the mel format is fine).

Now let me start v3c with 2x downsample:


python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3c --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.01 --recon_weight 2.0 --spectral_weight 1.5 --downsample_factor 2 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10 > /home/ubuntu/lewm-tts/train_v3c.log 2>&1

Start v3c training with 2x downsample


Command running in background with ID: bnrqjy832. Output is being written to: /tmp/claude-1001/-home-ubuntu/5610b655-b085-41d5-b3d9-c0750b48e28e/tasks/bnrqjy832.output


sleep 60 && head -10 /home/ubuntu/lewm-tts/train_v3c.log

Check v3c first epoch


Device: cuda
LeWM TTS model: 21.11M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=11.9772 pred=1.4899 kl=0.2497 recon=2.3543 spec=3.8524 | 42.9s
  New best model (loss=11.9772)


v3c running — **42.9s/epoch** (vs 27s for v3b, as expected from longer sequences). Total ~2.4 hours for 200 epochs.

Mel format is confirmed correct. The fundamental issue is the decoder's 4x compression — v3c with 2x should let us know if that's truly the bottleneck. If v3c still sounds robotic, then we need to rethink the decoder architecture entirely.


That's just the v3b process we killed (exit code 144 = SIGKILL). Expected. v3c is running fine.


Check v3c progress


Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=11.9772 pred=1.4899 kl=0.2497 recon=2.3543 spec=3.8524 | 42.9s
  New best model (loss=11.9772)


v3c is training. ~43s/epoch, will take ~2.4 hours for 200 epochs.


check logs


Check v3c training progress


Epoch 4/200 | avg_loss=10.6588 pred=0.9144 kl=0.2344 recon=2.3339 spec=3.3843 | 40.8s
  New best model (loss=10.6588)
Epoch 5/200 | avg_loss=9.0671 pred=0.8269 kl=0.3122 recon=2.2122 spec=2.5436 | 40.3s
  New best model (loss=9.0671)
  [6/200] step 150 | loss=7.2285 pred=0.8280 kl=0.2964 recon=1.6269 spec=2.0975 | grad=7.31 | 22.7s
Epoch 6/200 | avg_loss=7.2588 pred=0.7820 kl=0.2954 recon=1.6706 spec=2.0902 | 40.6s
  New best model (loss=7.2588)
Epoch 7/200 | avg_loss=5.9076 pred=0.6632 kl=0.3169 recon=1.3503 spec=1.6955 | 40.6s
  New best model (loss=5.9076)
  [8/200] step 200 | loss=4.9208 pred=0.5572 kl=0.3621 recon=1.1829 spec=1.3314 | grad=4.71 | 16.9s
Epoch 8/200 | avg_loss=4.8213 pred=0.5439 kl=0.3731 recon=1.1619 spec=1.3019 | 40.8s
  New best model (loss=4.8213)
Epoch 9/200 | avg_loss=4.0268 pred=0.4353 kl=0.4836 recon=1.0010 spec=1.0590 | 40.4s
  New best model (loss=4.0268)
  [10/200] step 250 | loss=3.8342 pred=0.3734 kl=0.5448 recon=0.9780 spec=1.0024 | grad=32.58 | 10.9s
Epoch 10/200 | avg_loss=3.7727 pred=0.3652 kl=0.5560 recon=0.9683 spec=0.9797 | 40.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3c/checkpoints/epoch_0010.pt
  New best model (loss=3.7727)
Epoch 11/200 | avg_loss=3.6309 pred=0.3330 kl=0.5902 recon=0.9486 spec=0.9328 | 40.5s
  New best model (loss=3.6309)


Epoch 11, recon=0.95. Comparing at same epoch with v3b:

| Epoch 10 | v3b (4x) | v3c (2x) |
|----------|----------|----------|
| recon | 0.96 | 0.97 |
| spectral | 0.90 | 0.98 |
| pred | 0.34 | 0.37 |

Very similar so far — early epochs are dominated by the decoder learning the overall mel distribution, not fine detail. The 2x advantage should show up later when it needs to reconstruct sharp features. Let's check back in ~30 min.


run a sample


python3 -c "
import torch, glob, torchaudio, json
from model import LeWMTTS
from vocos import Vocos

ckpt = torch.load('output_v3c/checkpoints/best.pt', map_location='cuda', weights_only=False)
model = LeWMTTS(ckpt['config']).cuda()
model.load_state_dict(ckpt['model'])
model.eval()
print(f'v3c epoch {ckpt[\"epoch\"]}')

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

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))
with open('processed_data_100mel/manifest.json') as f:
    manifest = json.load(f)

with torch.no_grad():
    for idx in [0, 100, 500]:
        mel = torch.load(mel_files[idx], weights_only=True).unsqueeze(0).cuda()
        z, mu, logvar = model.audio_encoder(mel)
        mel_recon = model.mel_decoder(mu)
        wav = vocos.decode(mel_recon)
        torchaudio.save(f'sample_recon_v3c_{idx}.wav', wav.cpu(), 24000)
        T = min(mel.shape[2], mel_recon.shape[2])
        l1 = (mel[:,:,:T] - mel_recon[:,:,:T]).abs().mean().item()
        print(f'Recon {idx}: L1={l1:.4f}')

    # TF
    mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0).cuda()
    text = manifest[0]['text']
    text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()
    z, mu, logvar = model.audio_encoder(mel)
    text_emb = model.text_encoder(text_tokens)
    predicted = model.predictor(mu[:, :-1], text_emb)
    combined = torch.cat([mu[:, :1], predicted], dim=1)
    mel_pred = model.mel_decoder(combined)
    wav = vocos.decode(mel_pred)
    torchaudio.save('sample_tf_v3c_0.wav', wav.cpu(), 24000)
    print('TF saved')

print('Done!')
" 2>&1

Generate v3c samples at current epoch


Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 6, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1530, in load
    return _load(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2122, in _load
    result = unpickler.load()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2086, in persistent_load
    typed_storage = load_tensor(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2052, in load_tensor
    wrap_storage = restore_location(storage, location)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1859, in restore_location
    return default_restore_location(storage, map_location)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 698, in default_restore_location
    result = fn(storage, location)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 637, in _deserialize
    return obj.to(device=device)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/storage.py", line 291, in to
    return _to(self, device, non_blocking)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/_utils.py", line 101, in _to
    untyped_storage = torch.UntypedStorage(self.size(), device=device)
torch.AcceleratorError: CUDA error: out of memory
Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


Training is using the GPU. Let me load on CPU instead:


python3 -c "
import torch, glob, torchaudio, json
from model import LeWMTTS
from vocos import Vocos

ckpt = torch.load('output_v3c/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTS(ckpt['config'])
model.load_state_dict(ckpt['model'])
model.eval()
print(f'v3c epoch {ckpt[\"epoch\"]}')

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

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))
with open('processed_data_100mel/manifest.json') as f:
    manifest = json.load(f)

with torch.no_grad():
    for idx in [0, 100, 500]:
        mel = torch.load(mel_files[idx], weights_only=True).unsqueeze(0)
        z, mu, logvar = model.audio_encoder(mel)
        mel_recon = model.mel_decoder(mu)
        wav = vocos.decode(mel_recon)
        torchaudio.save(f'sample_recon_v3c_{idx}.wav', wav, 24000)
        T = min(mel.shape[2], mel_recon.shape[2])
        l1 = (mel[:,:,:T] - mel_recon[:,:,:T]).abs().mean().item()
        print(f'Recon {idx}: L1={l1:.4f}')

    # TF
    mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0)
    text = manifest[0]['text']
    text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0)
    z, mu, logvar = model.audio_encoder(mel)
    text_emb = model.text_encoder(text_tokens)
    predicted = model.predictor(mu[:, :-1], text_emb)
    combined = torch.cat([mu[:, :1], predicted], dim=1)
    mel_pred = model.mel_decoder(combined)
    wav = vocos.decode(mel_pred)
    torchaudio.save('sample_tf_v3c_0.wav', wav, 24000)
    print('TF saved')

print('Done!')
" 2>&1

Generate v3c samples on CPU


v3c epoch 12
Recon 0: L1=0.8239
Recon 100: L1=0.8503
Recon 500: L1=0.8303
TF saved
Done!


Epoch 12 — still very early. L1=0.83 (v3b was 0.96 at epoch 10).

Samples in `/home/ubuntu/lewm-tts/`:
- `sample_recon_v3c_0/100/500.wav`
- `sample_tf_v3c_0.wav`

It's only epoch 12, so don't expect much yet. How does it compare to v3b at the same stage?


test again


Check current epoch


  New best model (loss=2.5070)
  [39/200] step 1050 | loss=2.5116 pred=0.0256 kl=2.3389 recon=0.6930 spec=0.7182 | grad=8.98 | 36.3s
Epoch 39/200 | avg_loss=2.5086 pred=0.0256 kl=2.3421 recon=0.6875 spec=0.7236 | 40.8s


python3 -c "
import torch, glob, torchaudio, json
from model import LeWMTTS
from vocos import Vocos

ckpt = torch.load('output_v3c/checkpoints/best.pt', map_location='cpu', weights_only=False)
model = LeWMTTS(ckpt['config'])
model.load_state_dict(ckpt['model'])
model.eval()
print(f'v3c epoch {ckpt[\"epoch\"]}')

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

mel_files = sorted(glob.glob('processed_data_100mel/mels/*.pt'))
with open('processed_data_100mel/manifest.json') as f:
    manifest = json.load(f)

with torch.no_grad():
    for idx in [0, 100, 500]:
        mel = torch.load(mel_files[idx], weights_only=True).unsqueeze(0)
        z, mu, logvar = model.audio_encoder(mel)
        mel_recon = model.mel_decoder(mu)
        wav = vocos.decode(mel_recon)
        torchaudio.save(f'sample_recon_v3c_{idx}.wav', wav, 24000)
        T = min(mel.shape[2], mel_recon.shape[2])
        l1 = (mel[:,:,:T] - mel_recon[:,:,:T]).abs().mean().item()
        print(f'Recon {idx}: L1={l1:.4f}')

    # TF
    mel = torch.load(mel_files[0], weights_only=True).unsqueeze(0)
    text = manifest[0]['text']
    text_tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0)
    z, mu, logvar = model.audio_encoder(mel)
    text_emb = model.text_encoder(text_tokens)
    predicted = model.predictor(mu[:, :-1], text_emb)
    combined = torch.cat([mu[:, :1], predicted], dim=1)
    mel_pred = model.mel_decoder(combined)
    wav = vocos.decode(mel_pred)
    torchaudio.save('sample_tf_v3c_0.wav', wav, 24000)
    print('TF saved')

    # AR prompted
    prompt_len = 40
    all_embs = [mu[:, :prompt_len]]
    cache = model.init_ar_cache()
    for i in range(prompt_len):
        _, cache = model.predict_next_cached(mu[:, i:i+1], text_emb, i, cache)
    emb = mu[:, prompt_len-1:prompt_len]
    for step in range(prompt_len, 400):
        emb, cache = model.predict_next_cached(emb, text_emb, step, cache)
        all_embs.append(emb)
    mel_ar = model.mel_decoder(torch.cat(all_embs, dim=1))
    wav_ar = vocos.decode(mel_ar)
    torchaudio.save('sample_ar_v3c.wav', wav_ar, 24000)
    print('AR prompted saved')

print('Done!')
" 2>&1

Generate v3c samples at epoch ~39