=== 1. Ground truth + Vocos ===
Saved: vocos_groundtruth.wav (9.51s)
=== 2. Reconstruction + Vocos ===
Saved: vocos_reconstruct.wav (9.50s)
=== 3. Teacher-forced predict + Vocos ===
Saved: vocos_predict_tf.wav (9.50s)

Done! Listen to vocos_groundtruth.wav first — if that sounds good, the vocoder is fixed.


All 3 files generated. Here's what to listen to in `/home/ubuntu/lewm-tts/output/`:

| File | What it tests |
|------|--------------|
| `vocos_groundtruth.wav` | Real mel → Vocos (should sound natural now) |
| `vocos_reconstruct.wav` | Encoder → Decoder → Vocos (tests model's mel quality) |
| `vocos_predict_tf.wav` | Encoder → Predictor → Decoder → Vocos (tests prediction) |

**Listen to `vocos_groundtruth.wav` first.** If that sounds clear and natural, then Griffin-Lim was the whole problem and switching to Vocos fixes the vocoder bottleneck. The remaining quality gap between groundtruth and reconstruct/predict will close as training continues (still at ~7%, 186 epochs to go).

I didn't generate AR synthesis with Vocos because it's pure noise at epoch 10 — that needs more training to work. Training is still running.


vocos_groundtruth.wav - good ..vocos_reconstruct.wav - robotic.. and other one too robotic


So the diagnosis is now clear:

1. **Vocos vocoder works great** — ground truth sounds good
2. **The mel decoder is the problem** — reconstruction is robotic, meaning the encoder→decoder path is losing quality
3. **Predictor is OK** — similar quality to reconstruction, so the predictor isn't making it worse

The issue is that the mel decoder (simple transposed convs) isn't reconstructing mels well enough. The L1 error was 0.59 which is still too high.

Options to fix:
1. **Better mel decoder architecture** — more layers, skip connections, residual blocks
2. **More reconstruction loss weight** — currently recon_weight=1.0, could increase

3. **Train longer** — only 10 epochs so far
4. **Switch to 100 mels** to match Vocos natively (avoid interpolation artifacts)

The most impactful fix is probably upgrading the mel decoder itself — the current two transposed convolutions are too simplistic. I should add residual blocks to make it more expressive. But I'm also wondering if the 80→100 mel interpolation we're doing for Vocos is introducing artifacts that are masking the real decoder quality. I should test the reconstruction against Griffin-Lim as well to isolate whether the problem is the decoder or the mel interpolation.

Actually, the user already confirmed Griffin-Lim sounds robotic too, so the mel decoder is definitely the bottleneck. Looking at the reconstruction loss of 0.68 L1 on log-mel — that's roughly 5% of the range given the values span from -8 to 4.5, but for good quality I probably need to get below 0.2. My priorities are: improve the decoder architecture first since that doesn't require reprocessing, maybe increase the reconstruction weight, and then switch to 100 mels natively if needed.

But I need to respect that the user said not to stop training. I can update the model architecture and prepare for the next run without interrupting the current one — since it's running from an old checkpoint directory, I can modify model.py and set up the improvements for when training resumes. Let me beef up the MelDecoder with residual blocks and plan the reprocessing with 100 mels for the next training iteration.


Good — that tells us exactly what's wrong. **Vocos is fine, the mel decoder is the bottleneck.** The current decoder is too simple (just 2 transposed convs). The mels coming out of it aren't detailed enough.

Two fixes needed:
1. **Better mel decoder** — add residual blocks for detail
2. **Switch to 100 mels** natively (avoid 80→100 interpolation artifacts)

Training is still running — I won't touch it. I'll prepare the improved version for the next run.


Bash

Check training progress


  New best model (loss=0.7040)
  [43/200] step 4540 | loss=0.7065 pred=0.0490 kl=1.2003 recon=0.5375 | grad=1.46 | 1.0s
  [43/200] step 4560 | loss=0.7029 pred=0.0490 kl=1.2016 recon=0.5337 | grad=1.21 | 4.9s


Good, epoch 43/200, recon loss down to 0.53. Now let me build the improved version.


Create v2 directory for improved model


(Bash completed with no output)


Write

"""
LeWM TTS v2 — Improved mel decoder + 100 mel bins for native Vocos support.

Changes from v1:
- MelDecoder: ResNet-style with skip connections, much more capacity
- 100 mel bins (matches Vocos 24kHz natively, no interpolation)
- PostNet for mel refinement
"""

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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):
    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

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

        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)

        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):
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)

        if mel_mask is not None:
            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)

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

        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)

        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

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

        predicted = self.output_proj(predicted)
        return predicted

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

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

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

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

class ResBlock(nn.Module):
    """Residual block with dilated convolutions."""
    def __init__(self, channels, kernel_size=3, dilation=1):
        super().__init__()
        padding = (kernel_size * dilation - dilation) // 2
        self.conv = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=padding, dilation=dilation),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
        )

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

class MelDecoder(nn.Module):
    """
    Improved mel decoder with:
    - Transposed conv upsampling (4x)
    - Residual blocks with dilated convolutions for detail
    - PostNet for mel refinement
    """

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        ch = d_model * 2  # 512 channels internally for more capacity

        self.pre = nn.Sequential(
            nn.Linear(d_model, ch),
            nn.GELU(),
        )

        # Upsample 4x with residual refinement at each stage
        self.up1 = nn.ConvTranspose1d(ch, ch, kernel_size=4, stride=2, padding=1)
        self.res1 = nn.Sequential(
            ResBlock(ch, kernel_size=3, dilation=1),
            ResBlock(ch, kernel_size=3, dilation=3),
            ResBlock(ch, kernel_size=3, dilation=5),
        )

        self.up2 = nn.ConvTranspose1d(ch, ch // 2, kernel_size=4, stride=2, padding=1)
        self.res2 = nn.Sequential(
            ResBlock(ch // 2, kernel_size=3, dilation=1),
            ResBlock(ch // 2, kernel_size=3, dilation=3),
            ResBlock(ch // 2, kernel_size=3, dilation=5),
        )

        # Project to mel
        self.mel_proj = nn.Conv1d(ch // 2, n_mels, kernel_size=7, padding=3)

        # PostNet: 5-layer CNN that refines mel output
        self.postnet = nn.Sequential(
            nn.Conv1d(n_mels, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, n_mels, kernel_size=5, padding=2),
        )

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

        x = F.gelu(self.up1(x))
        x = self.res1(x)

        x = F.gelu(self.up2(x))
        x = self.res2(x)

        mel_coarse = self.mel_proj(x)
        mel_fine = mel_coarse + self.postnet(mel_coarse)

        return mel_fine

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

class LeWMTTS(nn.Module):
    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.1)
        self.recon_weight = config.get("recon_weight", 2.0)

        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,
        )
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
        text_emb = self.text_encoder(text_tokens, text_mask)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        input_emb = z[:, :-1]
        target_emb = z[:, 1:]

        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
        if pred_mask is not None:
            valid = (~pred_mask).unsqueeze(-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
        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 (with PostNet — both coarse and fine)
        mel_recon = self.mel_decoder(z)
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        mel_target = mel
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel_target = 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()
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel_target * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel_target.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel_target)

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

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

    def encode_audio(self, mel):
        z, mu, _ = self.audio_encoder(mel)
        return mu

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]

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.1,
            "recon_weight": 2.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS v2: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


File created successfully at: /home/ubuntu/lewm-tts/v2/model.py

"""
LeWM TTS v2 — Improved mel decoder + 100 mel bins for native Vocos support.

Changes from v1:
- MelDecoder: ResNet-style with skip connections, much more capacity
- 100 mel bins (matches Vocos 24kHz natively, no interpolation)
- PostNet for mel refinement
"""

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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):
    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

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

        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)

        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):
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)

        if mel_mask is not None:
            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)

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

        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)

        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

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

        predicted = self.output_proj(predicted)
        return predicted

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

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

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

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

class ResBlock(nn.Module):
    """Residual block with dilated convolutions."""
    def __init__(self, channels, kernel_size=3, dilation=1):
        super().__init__()
        padding = (kernel_size * dilation - dilation) // 2
        self.conv = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=padding, dilation=dilation),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
        )

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

class MelDecoder(nn.Module):
    """
    Improved mel decoder with:
    - Transposed conv upsampling (4x)
    - Residual blocks with dilated convolutions for detail
    - PostNet for mel refinement
    """

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        ch = d_model * 2  # 512 channels internally for more capacity

        self.pre = nn.Sequential(
            nn.Linear(d_model, ch),
            nn.GELU(),
        )

        # Upsample 4x with residual refinement at each stage
        self.up1 = nn.ConvTranspose1d(ch, ch, kernel_size=4, stride=2, padding=1)
        self.res1 = nn.Sequential(
            ResBlock(ch, kernel_size=3, dilation=1),
            ResBlock(ch, kernel_size=3, dilation=3),
            ResBlock(ch, kernel_size=3, dilation=5),
        )

        self.up2 = nn.ConvTranspose1d(ch, ch // 2, kernel_size=4, stride=2, padding=1)
        self.res2 = nn.Sequential(
            ResBlock(ch // 2, kernel_size=3, dilation=1),
            ResBlock(ch // 2, kernel_size=3, dilation=3),
            ResBlock(ch // 2, kernel_size=3, dilation=5),
        )

        # Project to mel
        self.mel_proj = nn.Conv1d(ch // 2, n_mels, kernel_size=7, padding=3)

        # PostNet: 5-layer CNN that refines mel output
        self.postnet = nn.Sequential(
            nn.Conv1d(n_mels, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, ch // 2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch // 2),
            nn.Tanh(),
            nn.Conv1d(ch // 2, n_mels, kernel_size=5, padding=2),
        )

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

        x = F.gelu(self.up1(x))
        x = self.res1(x)

        x = F.gelu(self.up2(x))
        x = self.res2(x)

        mel_coarse = self.mel_proj(x)
        mel_fine = mel_coarse + self.postnet(mel_coarse)

        return mel_fine

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

class LeWMTTS(nn.Module):
    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.1)
        self.recon_weight = config.get("recon_weight", 2.0)

        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,
        )
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
        text_emb = self.text_encoder(text_tokens, text_mask)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        input_emb = z[:, :-1]
        target_emb = z[:, 1:]

        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
        if pred_mask is not None:
            valid = (~pred_mask).unsqueeze(-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
        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 (with PostNet — both coarse and fine)
        mel_recon = self.mel_decoder(z)
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        mel_target = mel
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel_target = 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()
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel_target * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel_target.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel_target)

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

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

    def encode_audio(self, mel):
        z, mu, _ = self.audio_encoder(mel)
        return mu

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]

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.1,
            "recon_weight": 2.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS v2: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


Now reprocess data with 100 mels and copy over the other files.


"""
Data preprocessing for LeWM TTS v2.
Same as v1 but with 100 mel bins to match Vocos 24kHz natively.
"""

import os
import json
import numpy as np
import torch
import torchaudio
from datasets import load_from_disk
from pathlib import Path
import argparse

MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 100,
    "f_min": 0,
    "f_max": 12000,
    "power": 1.0,
}

SEG_CONFIG = {
    "min_duration": 3.0,
    "max_duration": 12.0,
    "target_duration": 8.0,
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,
}

def detect_silence_points(waveform, sr, threshold_db=-40, min_silence_len=0.3):
    frame_length = int(0.025 * sr)
    hop = int(0.010 * sr)
    n_frames = 1 + (len(waveform) - frame_length) // hop
    rms_db = []
    for i in range(n_frames):
        start = i * hop
        frame = waveform[start:start + frame_length]
        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
        rms_db.append(20 * np.log10(rms + 1e-10))

    rms_db = np.array(rms_db)
    is_silence = rms_db < threshold_db
    silence_points = []
    min_frames = int(min_silence_len / 0.010)

    in_silence = False
    silence_start = 0
    for i, s in enumerate(is_silence):
        if s and not in_silence:
            silence_start = i
            in_silence = True
        elif not s and in_silence:
            if i - silence_start >= min_frames:
                mid_frame = (silence_start + i) // 2
                mid_sample = mid_frame * hop
                silence_points.append(mid_sample)
            in_silence = False

    return silence_points

def segment_audio(waveform, sr, text, config):
    total_duration = len(waveform) / sr
    if total_duration <= config["max_duration"]:
        if total_duration >= config["min_duration"]:
            return [(waveform, text)]
        else:
            return []

    silence_points = detect_silence_points(
        waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"]
    )

    split_points = [0] + silence_points + [len(waveform)]
    segments = []
    current_start = 0

    for i, point in enumerate(split_points[1:], 1):
        current_duration = (point - current_start) / sr
        if current_duration >= config["target_duration"]:
            chunk = waveform[current_start:point]
            chunk_duration = len(chunk) / sr
            if config["min_duration"] <= chunk_duration <= config["max_duration"] * 1.2:
                segments.append(chunk)
            elif chunk_duration > config["max_duration"] * 1.2:
                max_samples = int(config["max_duration"] * sr)
                for j in range(0, len(chunk), max_samples):
                    sub = chunk[j:j + max_samples]
                    if len(sub) / sr >= config["min_duration"]:
                        segments.append(sub)
            current_start = point

    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined

    results = []
    total_samples = sum(len(s) for s in segments)
    char_idx = 0
    for seg in segments:
        proportion = len(seg) / total_samples
        n_chars = max(1, int(len(text) * proportion))
        seg_text = text[char_idx:char_idx + n_chars]
        char_idx += n_chars
        results.append((seg, seg_text.strip()))

    return results

def extract_mel(waveform, config):
    mel_transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=config["sample_rate"],
        n_fft=config["n_fft"],
        hop_length=config["hop_length"],
        n_mels=config["n_mels"],
        f_min=config["f_min"],
        f_max=config["f_max"],
        power=config["power"],
    )
    mel = mel_transform(waveform)
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def trim_silence(waveform, sr, threshold_db=-45):
    threshold = 10 ** (threshold_db / 20)
    abs_wav = np.abs(waveform)
    window = int(0.01 * sr)
    start = 0
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break
    end = len(abs_wav)
    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break
    pad = int(0.05 * sr)
    start = max(0, start - pad)
    end = min(len(waveform), end + pad)
    return waveform[start:end]

def preprocess_dataset(input_path, output_path):
    print(f"Loading dataset from {input_path}...")
    ds = load_from_disk(input_path)

    output_dir = Path(output_path)
    audio_dir = output_dir / "audio"
    mel_dir = output_dir / "mels"
    audio_dir.mkdir(parents=True, exist_ok=True)
    mel_dir.mkdir(parents=True, exist_ok=True)

    manifest = []
    total_duration = 0.0
    idx = 0

    for i in range(len(ds)):
        sample = ds[i]
        audio_data = sample["audio"]
        waveform = np.array(audio_data["array"], dtype=np.float32)
        sr = audio_data["sampling_rate"]
        text = sample["text"]

        waveform = trim_silence(waveform, sr)
        segments = segment_audio(waveform, sr, text, SEG_CONFIG)

        for seg_wav, seg_text in segments:
            if not seg_text:
                continue
            duration = len(seg_wav) / sr
            total_duration += duration

            audio_path = audio_dir / f"{idx:06d}.wav"
            wav_tensor = torch.from_numpy(seg_wav).unsqueeze(0)
            torchaudio.save(str(audio_path), wav_tensor, sr)

            mel = extract_mel(wav_tensor, MEL_CONFIG)
            mel_path = mel_dir / f"{idx:06d}.pt"
            torch.save(mel.squeeze(0), str(mel_path))

            manifest.append({
                "id": idx,
                "audio_path": str(audio_path),
                "mel_path": str(mel_path),
                "text": seg_text,
                "duration": round(duration, 3),
                "mel_frames": mel.shape[-1],
            })
            idx += 1

        if (i + 1) % 100 == 0:
            print(f"  Processed {i+1}/{len(ds)} samples, {idx} segments, {total_duration/3600:.2f}h total")

    with open(output_dir / "manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)

    config = {"mel": MEL_CONFIG, "segmentation": SEG_CONFIG,
              "total_segments": idx, "total_duration_hours": round(total_duration / 3600, 2),
              "original_samples": len(ds)}
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nDone! {idx} segments, {total_duration/3600:.2f} hours")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="/home/ubuntu/modi_dataset")
    parser.add_argument("--output", default="/home/ubuntu/lewm-tts/v2/processed_data")
    args = parser.parse_args()
    preprocess_dataset(args.input, args.output)


File created successfully at: /home/ubuntu/lewm-tts/v2/preprocess.py

"""
Data preprocessing for LeWM TTS v2.
Same as v1 but with 100 mel bins to match Vocos 24kHz natively.
"""

import os
import json
import numpy as np
import torch
import torchaudio
from datasets import load_from_disk
from pathlib import Path
import argparse

MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 100,
    "f_min": 0,
    "f_max": 12000,
    "power": 1.0,
}

SEG_CONFIG = {
    "min_duration": 3.0,
    "max_duration": 12.0,
    "target_duration": 8.0,
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,
}

def detect_silence_points(waveform, sr, threshold_db=-40, min_silence_len=0.3):
    frame_length = int(0.025 * sr)
    hop = int(0.010 * sr)
    n_frames = 1 + (len(waveform) - frame_length) // hop
    rms_db = []
    for i in range(n_frames):
        start = i * hop
        frame = waveform[start:start + frame_length]
        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
        rms_db.append(20 * np.log10(rms + 1e-10))

    rms_db = np.array(rms_db)
    is_silence = rms_db < threshold_db
    silence_points = []
    min_frames = int(min_silence_len / 0.010)

    in_silence = False
    silence_start = 0
    for i, s in enumerate(is_silence):
        if s and not in_silence:
            silence_start = i
            in_silence = True
        elif not s and in_silence:
            if i - silence_start >= min_frames:
                mid_frame = (silence_start + i) // 2
                mid_sample = mid_frame * hop
                silence_points.append(mid_sample)
            in_silence = False

    return silence_points

def segment_audio(waveform, sr, text, config):
    total_duration = len(waveform) / sr
    if total_duration <= config["max_duration"]:
        if total_duration >= config["min_duration"]:
            return [(waveform, text)]
        else:
            return []

    silence_points = detect_silence_points(
        waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"]
    )

    split_points = [0] + silence_points + [len(waveform)]
    segments = []
    current_start = 0

    for i, point in enumerate(split_points[1:], 1):
        current_duration = (point - current_start) / sr
        if current_duration >= config["target_duration"]:
            chunk = waveform[current_start:point]
            chunk_duration = len(chunk) / sr
            if config["min_duration"] <= chunk_duration <= config["max_duration"] * 1.2:
                segments.append(chunk)
            elif chunk_duration > config["max_duration"] * 1.2:
                max_samples = int(config["max_duration"] * sr)
                for j in range(0, len(chunk), max_samples):
                    sub = chunk[j:j + max_samples]
                    if len(sub) / sr >= config["min_duration"]:
                        segments.append(sub)
            current_start = point

    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined

    results = []
    total_samples = sum(len(s) for s in segments)
    char_idx = 0
    for seg in segments:
        proportion = len(seg) / total_samples
        n_chars = max(1, int(len(text) * proportion))
        seg_text = text[char_idx:char_idx + n_chars]
        char_idx += n_chars
        results.append((seg, seg_text.strip()))

    return results

def extract_mel(waveform, config):
    mel_transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=config["sample_rate"],
        n_fft=config["n_fft"],
        hop_length=config["hop_length"],
        n_mels=config["n_mels"],
        f_min=config["f_min"],
        f_max=config["f_max"],
        power=config["power"],
    )
    mel = mel_transform(waveform)
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def trim_silence(waveform, sr, threshold_db=-45):
    threshold = 10 ** (threshold_db / 20)
    abs_wav = np.abs(waveform)
    window = int(0.01 * sr)
    start = 0
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break
    end = len(abs_wav)
    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break
    pad = int(0.05 * sr)
    start = max(0, start - pad)
    end = min(len(waveform), end + pad)
    return waveform[start:end]

def preprocess_dataset(input_path, output_path):
    print(f"Loading dataset from {input_path}...")
    ds = load_from_disk(input_path)

    output_dir = Path(output_path)
    audio_dir = output_dir / "audio"
    mel_dir = output_dir / "mels"
    audio_dir.mkdir(parents=True, exist_ok=True)
    mel_dir.mkdir(parents=True, exist_ok=True)

    manifest = []
    total_duration = 0.0
    idx = 0

    for i in range(len(ds)):
        sample = ds[i]
        audio_data = sample["audio"]
        waveform = np.array(audio_data["array"], dtype=np.float32)
        sr = audio_data["sampling_rate"]
        text = sample["text"]

        waveform = trim_silence(waveform, sr)
        segments = segment_audio(waveform, sr, text, SEG_CONFIG)

        for seg_wav, seg_text in segments:
            if not seg_text:
                continue
            duration = len(seg_wav) / sr
            total_duration += duration

            audio_path = audio_dir / f"{idx:06d}.wav"
            wav_tensor = torch.from_numpy(seg_wav).unsqueeze(0)
            torchaudio.save(str(audio_path), wav_tensor, sr)

            mel = extract_mel(wav_tensor, MEL_CONFIG)
            mel_path = mel_dir / f"{idx:06d}.pt"
            torch.save(mel.squeeze(0), str(mel_path))

            manifest.append({
                "id": idx,
                "audio_path": str(audio_path),
                "mel_path": str(mel_path),
                "text": seg_text,
                "duration": round(duration, 3),
                "mel_frames": mel.shape[-1],
            })
            idx += 1

        if (i + 1) % 100 == 0:
            print(f"  Processed {i+1}/{len(ds)} samples, {idx} segments, {total_duration/3600:.2f}h total")

    with open(output_dir / "manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)

    config = {"mel": MEL_CONFIG, "segmentation": SEG_CONFIG,
              "total_segments": idx, "total_duration_hours": round(total_duration / 3600, 2),
              "original_samples": len(ds)}
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nDone! {idx} segments, {total_duration/3600:.2f} hours")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="/home/ubuntu/modi_dataset")
    parser.add_argument("--output", default="/home/ubuntu/lewm-tts/v2/processed_data")
    args = parser.parse_args()
    preprocess_dataset(args.input, args.output)


Bash

cp /home/ubuntu/lewm-tts/dataset.py /home/ubuntu/lewm-tts/v2/dataset.py && cp /home/ubuntu/lewm-tts/train.py /home/ubuntu/lewm-tts/v2/train.py && cp /home/ubuntu/lewm-tts/inference.py /home/ubuntu/lewm-tts/v2/inference.py

Copy dataset, train, inference to v2


Now test the v2 model size and run the preprocessing.


cd /home/ubuntu/lewm-tts/v2 && python3 -c "
from model import build_model
import torch

model, config = build_model()
device = torch.device('cuda')
model = model.to(device)

# Test forward
B = 2
mel = torch.randn(B, 100, 400).to(device)
text = torch.randint(0, 256, (B, 50)).to(device)
mel_mask = torch.zeros(B, 400, dtype=torch.bool).to(device)
text_mask = torch.zeros(B, 50, dtype=torch.bool).to(device)

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

# Test mel decoder quality
model.eval()
with torch.no_grad():
    z, mu, _ = model.audio_encoder(mel)
    mel_recon = model.mel_decoder(mu)
    print(f'\\nEncoder: {mel.shape} -> {mu.shape}')
    print(f'Decoder: {mu.shape} -> {mel_recon.shape}')
" 2>&1

Test v2 model


LeWM TTS v2: 22.83M parameters
Forward/backward OK
  total_loss: 3.2450
  prediction_loss: 1.5193
  kl_loss: 0.2567
  recon_loss: 0.8500

Encoder: torch.Size([2, 100, 400]) -> torch.Size([2, 100, 256])
Decoder: torch.Size([2, 100, 256]) -> torch.Size([2, 100, 400])


22.83M params (up from 14.5M — most of the growth is in the mel decoder). Now preprocess with 100 mels.


Preprocess dataset with 100 mel bins


Error: Exit code 137
[Request interrupted by user for tool use]


[Request interrupted by user for tool use]


i ahve question y we not used DAC vae


The user is asking why we didn't use DAC (Descript Audio Codec) VAE instead of mel spectrograms. This is a great question.

DAC is a neural audio codec that compresses audio into discrete tokens using a VQ-VAE architecture. It's used by models like Bark, VALL-E, and others for high-quality audio tokenization.

Let me explain the tradeoffs.


Really good question. DAC (Descript Audio Codec) would actually be a better fit here. Here's the comparison:

**Why we used mel spectrograms:**
- Simpler to start with, well-understood
- No pretrained model dependency
- Easier to debug

**Why DAC would be better:**

| | Mel + Vocoder | DAC |
|--|--|--|
| **Audio quality** | Depends on vocoder (Griffin-Lim = robotic) | Built-in decoder, near-lossless at 24kHz |
| **Representation** | Continuous, needs separate vocoder | Self-contained encode/decode |
| **Latent space** | Raw spectral features | Learned, compressed, meaningful |
| **Reconstruction** | We're training a mel decoder from scratch (our bottleneck!) | Pretrained decoder already perfect |
| **Pipeline complexity** | Encoder + Decoder + Vocoder (3 things to get right) | Just DAC encode/decode (1 thing) |

**The key insight:** Our biggest problem right now is the mel decoder producing robotic audio. With DAC, **we don't need a mel decoder at all**. DAC's pretrained decoder already reconstructs perfect audio. Our model only needs to predict DAC latent embeddings — and DAC handles the rest.

**How it would work:**
1. Preprocess: audio → DAC encoder → continuous latents (before quantization)
2. Train: JEPA predictor learns to predict next DAC latent from text + context
3. Inference: predicted latents → DAC decoder → waveform

This eliminates the mel decoder entirely — which is exactly where quality is breaking down.

Want me to rebuild with DAC? The v1 training is still running and we can let it finish, but DAC would be the smarter v2.


what will get if we finish v1 .. do we knwo what is it like with taht ?


Honestly? Even if v1 finishes all 200 epochs:

**What you'll get:**
- Reconstruction will improve somewhat (recon loss is still dropping — was 0.68, now at 0.53)
- Teacher-forced prediction will sound similar to reconstruction
- AR synthesis will probably still be noisy or very distorted

**Why the ceiling is low:**
- The mel decoder is too simple (v1) — even with perfect latents, it can't reconstruct detailed mels
- Griffin-Lim is fundamentally robotic — it's a 1980s algorithm
- Even with Vocos (100-mel interpolation hack), we're working around problems, not solving them

**The core issue:** We're training a mel decoder from scratch with 14.5M params to do something that DAC's pretrained decoder already does perfectly with 0 additional training.

**My honest recommendation:**

Let v1 finish (it's running, ~1hr left, costs nothing). It proves the JEPA predictor works — prediction loss dropped 99%. That's the valuable part. But don't expect listenable speech from it.

For v2, switch to DAC. It eliminates:
- The mel decoder (our bottleneck)
- The vocoder problem (Griffin-Lim/Vocos)
- The 80 vs 100 mel hassle

The JEPA predictor + text encoder stay exactly the same — that code carries over. Only the audio representation changes.

Want me to start building the DAC version now while v1 finishes in the background?


yeah what will be params ?


pip install descript-audio-codec 2>&1 | tail -3

Install DAC


vllm 0.16.0 requires protobuf!=6.30.*,!=6.31.*,!=6.32.*,!=6.33.0.*,!=6.33.1.*,!=6.33.2.*,!=6.33.3.*,!=6.33.4.*,>=5.29.6, but you have protobuf 4.25.8 which is incompatible.
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Successfully installed protobuf-4.25.8


python3 -c "
import dac
import torch

# Load pretrained DAC 24kHz
model_path = dac.utils.download(model_type='24khz')
dac_model = dac.DAC.from_pretrained(model_path)
dac_model.eval()

print('=== DAC 24kHz Model ===')
print(f'Sample rate: 24000')
print(f'Hop length: {dac_model.hop_length}')
print(f'Encoder output dim: {dac_model.encoder_dim}')
print(f'Latent dim (before quantize): {dac_model.latent_dim}')
print(f'Number of codebooks: {dac_model.n_codebooks}')
print(f'Codebook size: {dac_model.codebook_size}')

# Test encode/decode
wav = torch.randn(1, 1, 24000 * 5)  # 5 seconds
dac_model = dac_model.to('cuda')
wav = wav.to('cuda')

with torch.no_grad():
    # Get continuous latents (before quantization)
    z = dac_model.encoder(wav)
    print(f'\\nEncoder output shape: {z.shape}')  # [B, latent_dim, T_compressed]
    print(f'Compression ratio: {wav.shape[-1]} samples -> {z.shape[-1]} frames = {wav.shape[-1]/z.shape[-1]:.0f}x')
    print(f'Frame rate: {24000 / (wav.shape[-1]/z.shape[-1]):.1f} Hz')

    # Full encode (with quantization)
    x = dac_model.preprocess(wav, 24000)
    z_q, codes, latents, commitment_loss, codebook_loss = dac_model.encode(x)
    print(f'\\nQuantized latent shape: {z_q.shape}')
    print(f'Codes shape: {codes.shape}')

    # Decode back
    wav_recon = dac_model.decode(z_q)
    print(f'Reconstructed wav: {wav_recon.shape}')

    # We can also use continuous latents (pre-quantization)
    z_pre = dac_model.encoder(x)
    z_proj = dac_model.quantizer.from_latents(z_pre)
    print(f'\\nPre-quantize latent: {z_pre.shape}')
" 2>&1

Inspect DAC model dimensions


Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 7, in <module>
AttributeError: type object 'DAC' has no attribute 'from_pretrained'


python3 -c "
import dac
print(dir(dac))
print(dir(dac.DAC))
" 2>&1

Check DAC API


['DAC', 'DACFile', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__model_version__', '__name__', '__package__', '__path__', '__spec__', '__version__', 'audiotools', 'model', 'nn', 'utils']
['EXTERN', 'INTERN', 'T_destination', '__annotations__', '__call__', '__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattr__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setstate__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_apply', '_call_impl', '_compiled_call_impl', '_get_backward_hooks', '_get_backward_pre_hooks', '_get_name', '_load_from_state_dict', '_load_package', '_maybe_warn_non_full_backward_hook', '_named_members', '_register_load_state_dict_pre_hook', '_register_state_dict_hook', '_replicate_for_data_parallel', '_save_package', '_save_to_state_dict', '_slow_forward', '_version', '_wrapped_call_impl', 'add_module', 'apply', 'bfloat16', 'buffers', 'call_super_init', 'children', 'compile', 'compress', 'cpu', 'cuda', 'decode', 'decompress', 'device', 'double', 'dump_patches', 'encode', 'eval', 'extra_repr', 'float', 'forward', 'get_buffer', 'get_delay', 'get_extra_state', 'get_output_length', 'get_parameter', 'get_submodule', 'half', 'ipu', 'load', 'load_from_folder', 'load_state_dict', 'modules', 'mtia', 'named_buffers', 'named_children', 'named_modules', 'named_parameters', 'padding', 'parameters', 'preprocess', 'register_backward_hook', 'register_buffer', 'register_forward_hook', 'register_forward_pre_hook', 'register_full_backward_hook', 'register_full_backward_pre_hook', 'register_load_state_dict_post_hook', 'register_load_state_dict_pre_hook', 'register_module', 'register_parameter', 'register_state_dict_post_hook', 'register_state_dict_pre_hook', 'requires_grad_', 'save', 'save_to_folder', 'set_extra_state', 'set_submodule', 'share_memory', 'state_dict', 'to', 'to_empty', 'train', 'type', 'xpu', 'zero_grad']


python3 -c "
import dac
import torch

# Load DAC 24kHz
model_path = dac.utils.download(model_type='24khz')
dac_model = dac.DAC.load(model_path)
dac_model.eval()
dac_model = dac_model.to('cuda')

print('=== DAC 24kHz ===')
print(f'Hop length: {dac_model.hop_length}')
print(f'Encoder dim: {dac_model.encoder_dim}')
print(f'Latent dim: {dac_model.latent_dim}')
print(f'Codebooks: {dac_model.n_codebooks}')
print(f'Codebook size: {dac_model.codebook_size}')

# Test with 5s audio
wav = torch.randn(1, 1, 24000 * 5).to('cuda')

with torch.no_grad():
    x = dac_model.preprocess(wav, 24000)

    # Get continuous latents from encoder
    z_enc = dac_model.encoder(x)
    print(f'\nEncoder output: {z_enc.shape}')

    # Full encode path
    z_q, codes, latents, _, _ = dac_model.encode(x)
    print(f'Quantized z: {z_q.shape}')
    print(f'Codes: {codes.shape}')
    print(f'Latents: {latents.shape}')

    compression = wav.shape[-1] / z_q.shape[-1]
    print(f'Compression: {compression:.0f}x ({24000/compression:.1f} Hz frame rate)')

    # Decode
    wav_out = dac_model.decode(z_q)
    print(f'Decoded wav: {wav_out.shape}')

    # IMPORTANT: What dim does our JEPA predictor need to work with?
    print(f'\n=== For LeWM TTS ===')
    print(f'DAC latent dim to predict: {z_q.shape[1]}')
    print(f'DAC temporal frames per second: {24000/compression:.1f}')
    print(f'For 10s audio: {int(10 * 24000/compression)} frames to predict')
" 2>&1

Get DAC dimensions for model planning


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
=== DAC 24kHz ===
Hop length: 320
Encoder dim: 64
Latent dim: 1024
Codebooks: 32
Codebook size: 1024

Encoder output: torch.Size([1, 1024, 375])
Quantized z: torch.Size([1, 1024, 375])
Codes: torch.Size([1, 32, 375])
Latents: torch.Size([1, 256, 375])
Compression: 320x (75.0 Hz frame rate)
Decoded wav: torch.Size([1, 1, 119992])

=== For LeWM TTS ===
DAC latent dim to predict: 1024
DAC temporal frames per second: 75.0
For 10s audio: 750 frames to predict


Now I have all the numbers. Here's the comparison:


[Request interrupted by user]


what will be total size if we use dac


python3 -c "
import torch.nn as nn

# DAC latent dim = 1024, but we can project it down
# Option A: Work with full 1024 latent dim (expensive)
# Option B: Project 1024 -> 256 (our d_model), predict in 256, project back

# Let's calculate both

def count(module):
    return sum(p.numel() for p in module.parameters())

d_model = 256
nhead = 4

# Text Encoder (same as before)
text_enc = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model, nhead, d_model*4, batch_first=True, activation='gelu'),
    num_layers=4
)
text_params = count(nn.Embedding(256, d_model)) + count(text_enc) + count(nn.Linear(d_model, d_model))

# DAC projections (new — replace AudioEncoder + MelDecoder)
dac_in_proj = count(nn.Linear(1024, d_model))   # DAC latent -> d_model
dac_out_proj = count(nn.Linear(d_model, 1024))   # d_model -> DAC latent

# Predictor (same as before)
predictor = nn.TransformerDecoder(
    nn.TransformerDecoderLayer(d_model, nhead, d_model*4, batch_first=True, activation='gelu'),
    num_layers=6
)
pred_params = count(nn.Linear(d_model, d_model)) + count(predictor) + count(nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, d_model)))

# mu/logvar projections for Gaussian regularizer
gauss_params = count(nn.Linear(d_model, d_model)) * 2

total_trainable = text_params + dac_in_proj + dac_out_proj + pred_params + gauss_params

print('=== LeWM TTS with DAC (v2) ===')
print(f'Text Encoder:     {text_params/1e6:.2f}M')
print(f'DAC in proj:      {dac_in_proj/1e6:.2f}M')
print(f'DAC out proj:     {dac_out_proj/1e6:.2f}M')
print(f'Gaussian projs:   {gauss_params/1e6:.2f}M')
print(f'JEPA Predictor:   {pred_params/1e6:.2f}M')
print(f'─────────────────────────')
print(f'Total trainable:  {total_trainable/1e6:.2f}M')
print()
print('DAC model (frozen, not trained): ~74M')
print(f'Total on GPU at inference: ~{(total_trainable + 74e6)/1e6:.0f}M')
print()

# Compare
print('=== Comparison ===')
print(f'v1 (mel + Griffin-Lim):   14.50M trainable')
print(f'v2 (DAC):                 {total_trainable/1e6:.2f}M trainable')
print(f'')
print(f'What we REMOVED:')
print(f'  AudioEncoder (CNN+Transformer): ~4M')
print(f'  MelDecoder: ~3M')
print(f'  = replaced by 2 linear projections: {(dac_in_proj+dac_out_proj)/1e6:.2f}M')
" 2>&1

Calculate DAC version model size


=== LeWM TTS with DAC (v2) ===
Text Encoder:     3.29M
DAC in proj:      0.26M
DAC out proj:     0.26M
Gaussian projs:   0.13M
JEPA Predictor:   6.52M
─────────────────────────
Total trainable:  10.47M

DAC model (frozen, not trained): ~74M
Total on GPU at inference: ~84M

=== Comparison ===
v1 (mel + Griffin-Lim):   14.50M trainable
v2 (DAC):                 10.47M trainable

What we REMOVED:
  AudioEncoder (CNN+Transformer): ~4M
  MelDecoder: ~3M
  = replaced by 2 linear projections: 0.53M


**~10.5M trainable parameters** (actually smaller than v1!)

| | v1 (mel) | v2 (DAC) |
|--|--|--|
| **Trainable params** | 14.50M | 10.47M |
| **Frozen params** | 0 | ~74M (DAC, pretrained) |
| **Total on GPU** | 14.50M | ~84M |
| **Audio quality ceiling** | Limited by mel decoder | Near-lossless (DAC) |
| **Training speed** | ~24s/epoch | Should be similar or faster |

The key win: we train **less** but get **better audio** because DAC's pretrained decoder does the hard part. Our model just needs to learn two linear projections (1024↔256) instead of an entire encoder + decoder from scratch.

GPU memory at batch 64 would go from ~7GB to maybe ~12-15GB — still well within the A100's 80GB.

Want me to build it?


yeah run lets see .. if it goes.. we may test with speaker labels on this 84m if we can get 10 to 20 speakers with this 80m .. we can do onnx and run on edge .. we can do this liek 1000 speakers alsofi neede taht wont be a problem ... run this first andlest see


The user wants me to build the DAC version, run it, and is already thinking ahead about multi-speaker scaling. Let me build it efficiently and get it running.

Let me:
1. Build the DAC-based model
2. Preprocess data with DAC
3. Train
4. All while v1 is still running in background


Let me build it and get it running.


TaskCreate

Build LeWM TTS v2 with DAC codec

Model, preprocessing, training, inference — all using DAC latents instead of mel spectrograms. 10.5M trainable params + 74M frozen DAC.


Task #6 created successfully: Build LeWM TTS v2 with DAC codec

Build LeWM TTS v2 with DAC codec


Bash

Clean v2 directory for fresh DAC build


Updated task #6 status


(Bash completed with no output)


Write

"""
LeWM TTS v2 — DAC-based JEPA Text-to-Speech

Architecture:
  - DAC encoder (frozen): audio → 1024-dim continuous latents @ 75Hz
  - Linear proj: 1024 → 256 (d_model)
  - TextEncoder: byte-level transformer
  - JEPAPredictor: causal transformer decoder with cross-attention to text
  - Linear proj: 256 → 1024 (back to DAC space)
  - DAC decoder (frozen): latents → waveform

  Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL)
  NO mel decoder needed — DAC handles audio reconstruction.
"""

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

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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)
        return self.proj(x)

class JEPAPredictor(nn.Module):
    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        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,
        )
        return self.output_proj(predicted)

class LeWMTTS(nn.Module):
    """
    LeWM TTS v2 with DAC latents.

    Training:
      - DAC encodes audio → 1024-dim latents (frozen)
      - Project to d_model, predict next embedding, KL regularize
    Inference:
      - Text → predict DAC latents autoregressively → DAC decode → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        dac_dim = config.get("dac_dim", 1024)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.1)

        # DAC space ↔ model space projections
        self.dac_in_proj = nn.Linear(dac_dim, d_model)
        self.dac_out_proj = nn.Linear(d_model, dac_dim)

        # Gaussian regularizer projections
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

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

        self.config = config

    def forward(self, dac_latents, text_tokens, latent_mask=None, text_mask=None):
        """
        Args:
            dac_latents: [B, T, 1024] — continuous DAC encoder output
            text_tokens: [B, T_text] — byte-level text tokens
            latent_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        """
        # Project DAC latents to model space
        h = self.dac_in_proj(dac_latents)  # [B, T, d_model]

        # Gaussian regularizer: project to mu/logvar
        mu = self.proj_mu(h)
        logvar = self.proj_logvar(h)

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

        # Encode text
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Predict next embedding
        input_emb = z[:, :-1]
        target_emb = z[:, 1:]

        if latent_mask is not None:
            pred_mask = latent_mask[:, :-1]
        else:
            pred_mask = None

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

        # Loss 1: Next-embedding prediction (MSE)
        if pred_mask is not None:
            valid = (~pred_mask).unsqueeze(-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: KL(N(mu, sigma²) || N(0, I))
        if latent_mask is not None:
            valid = (~latent_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())

        total_loss = prediction_loss + self.kl_weight * kl_loss

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

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """AR inference: predict next embedding from context."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]

    def latents_to_dac(self, z):
        """Project model-space embeddings back to DAC latent space."""
        return self.dac_out_proj(z)

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,
            "dac_dim": 1024,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.1,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS v2 (DAC): {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config


"""
LeWM TTS v2 — DAC-based JEPA Text-to-Speech

Architecture:
  - DAC encoder (frozen): audio → 1024-dim continuous latents @ 75Hz
  - Linear proj: 1024 → 256 (d_model)
  - TextEncoder: byte-level transformer
  - JEPAPredictor: causal transformer decoder with cross-attention to text
  - Linear proj: 256 → 1024 (back to DAC space)
  - DAC decoder (frozen): latents → waveform

  Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL)
  NO mel decoder needed — DAC handles audio reconstruction.
"""

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

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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)
        return self.proj(x)

class JEPAPredictor(nn.Module):
    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        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,
        )
        return self.output_proj(predicted)

class LeWMTTS(nn.Module):
    """
    LeWM TTS v2 with DAC latents.

    Training:
      - DAC encodes audio → 1024-dim latents (frozen)
      - Project to d_model, predict next embedding, KL regularize
    Inference:
      - Text → predict DAC latents autoregressively → DAC decode → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        dac_dim = config.get("dac_dim", 1024)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.1)

        # DAC space ↔ model space projections
        self.dac_in_proj = nn.Linear(dac_dim, d_model)
        self.dac_out_proj = nn.Linear(d_model, dac_dim)

        # Gaussian regularizer projections
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

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

        self.config = config

    def forward(self, dac_latents, text_tokens, latent_mask=None, text_mask=None):
        """
        Args:
            dac_latents: [B, T, 1024] — continuous DAC encoder output
            text_tokens: [B, T_text] — byte-level text tokens
            latent_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        """
        # Project DAC latents to model space
        h = self.dac_in_proj(dac_latents)  # [B, T, d_model]

        # Gaussian regularizer: project to mu/logvar
        mu = self.proj_mu(h)
        logvar = self.proj_logvar(h)

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

        # Encode text
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Predict next embedding
        input_emb = z[:, :-1]
        target_emb = z[:, 1:]

        if latent_mask is not None:
            pred_mask = latent_mask[:, :-1]
        else:
            pred_mask = None

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

        # Loss 1: Next-embedding prediction (MSE)
        if pred_mask is not None:
            valid = (~pred_mask).unsqueeze(-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: KL(N(mu, sigma²) || N(0, I))
        if latent_mask is not None:
            valid = (~latent_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())

        total_loss = prediction_loss + self.kl_weight * kl_loss

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

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """AR inference: predict next embedding from context."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]

    def latents_to_dac(self, z):
        """Project model-space embeddings back to DAC latent space."""
        return self.dac_out_proj(z)

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,
            "dac_dim": 1024,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.1,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS v2 (DAC): {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config


"""
Preprocess Modi dataset → DAC latents.
Encodes all audio through frozen DAC encoder, saves continuous latents.
"""

import os
import json
import numpy as np
import torch
import torchaudio
import dac
from datasets import load_from_disk
from pathlib import Path
import argparse

SEG_CONFIG = {
    "min_duration": 3.0,
    "max_duration": 12.0,
    "target_duration": 8.0,
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,
}

def detect_silence_points(waveform, sr, threshold_db=-40, min_silence_len=0.3):
    frame_length = int(0.025 * sr)
    hop = int(0.010 * sr)
    n_frames = 1 + (len(waveform) - frame_length) // hop
    rms_db = []
    for i in range(n_frames):
        start = i * hop
        frame = waveform[start:start + frame_length]
        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
        rms_db.append(20 * np.log10(rms + 1e-10))
    rms_db = np.array(rms_db)
    is_silence = rms_db < threshold_db
    silence_points = []
    min_frames = int(min_silence_len / 0.010)
    in_silence = False
    silence_start = 0
    for i, s in enumerate(is_silence):
        if s and not in_silence:
            silence_start = i
            in_silence = True
        elif not s and in_silence:
            if i - silence_start >= min_frames:
                mid_frame = (silence_start + i) // 2
                silence_points.append(mid_frame * hop)
            in_silence = False
    return silence_points

def segment_audio(waveform, sr, text, config):
    total_duration = len(waveform) / sr
    if total_duration <= config["max_duration"]:
        return [(waveform, text)] if total_duration >= config["min_duration"] else []
    silence_points = detect_silence_points(waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"])
    split_points = [0] + silence_points + [len(waveform)]
    segments = []
    current_start = 0
    for point in split_points[1:]:
        current_duration = (point - current_start) / sr
        if current_duration >= config["target_duration"]:
            chunk = waveform[current_start:point]
            chunk_duration = len(chunk) / sr
            if config["min_duration"] <= chunk_duration <= config["max_duration"] * 1.2:
                segments.append(chunk)
            elif chunk_duration > config["max_duration"] * 1.2:
                max_samples = int(config["max_duration"] * sr)
                for j in range(0, len(chunk), max_samples):
                    sub = chunk[j:j + max_samples]
                    if len(sub) / sr >= config["min_duration"]:
                        segments.append(sub)
            current_start = point
    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined
    results = []
    total_samples = sum(len(s) for s in segments)
    char_idx = 0
    for seg in segments:
        proportion = len(seg) / total_samples
        n_chars = max(1, int(len(text) * proportion))
        seg_text = text[char_idx:char_idx + n_chars].strip()
        char_idx += n_chars
        results.append((seg, seg_text))
    return results

def trim_silence(waveform, sr, threshold_db=-45):
    threshold = 10 ** (threshold_db / 20)
    abs_wav = np.abs(waveform)
    window = int(0.01 * sr)
    start = 0
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break
    end = len(abs_wav)
    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break
    pad = int(0.05 * sr)
    return waveform[max(0, start - pad):min(len(waveform), end + pad)]

def preprocess_dataset(input_path, output_path):
    print(f"Loading dataset from {input_path}...")
    ds = load_from_disk(input_path)

    # Load DAC model
    print("Loading DAC 24kHz model...")
    model_path = dac.utils.download(model_type='24khz')
    dac_model = dac.DAC.load(model_path)
    dac_model.eval()
    dac_model = dac_model.to('cuda')
    print(f"DAC loaded. Hop={dac_model.hop_length}, latent_dim={dac_model.latent_dim}")

    output_dir = Path(output_path)
    latent_dir = output_dir / "latents"
    audio_dir = output_dir / "audio"
    latent_dir.mkdir(parents=True, exist_ok=True)
    audio_dir.mkdir(parents=True, exist_ok=True)

    manifest = []
    total_duration = 0.0
    idx = 0

    for i in range(len(ds)):
        sample = ds[i]
        audio_data = sample["audio"]
        waveform = np.array(audio_data["array"], dtype=np.float32)
        sr = audio_data["sampling_rate"]
        text = sample["text"]

        waveform = trim_silence(waveform, sr)
        segments = segment_audio(waveform, sr, text, SEG_CONFIG)

        for seg_wav, seg_text in segments:
            if not seg_text:
                continue
            duration = len(seg_wav) / sr
            total_duration += duration

            # Save audio
            audio_path = audio_dir / f"{idx:06d}.wav"
            wav_tensor = torch.from_numpy(seg_wav).unsqueeze(0)
            torchaudio.save(str(audio_path), wav_tensor, sr)

            # Encode with DAC → continuous latents
            with torch.no_grad():
                wav_input = wav_tensor.unsqueeze(0).to('cuda')  # [1, 1, samples]
                wav_input = dac_model.preprocess(wav_input, sr)
                z = dac_model.encoder(wav_input)  # [1, 1024, T]
                z = z.squeeze(0).cpu()  # [1024, T]

            latent_path = latent_dir / f"{idx:06d}.pt"
            torch.save(z, str(latent_path))

            manifest.append({
                "id": idx,
                "audio_path": str(audio_path),
                "latent_path": str(latent_path),
                "text": seg_text,
                "duration": round(duration, 3),
                "latent_frames": z.shape[1],
            })
            idx += 1

        if (i + 1) % 100 == 0:
            print(f"  Processed {i+1}/{len(ds)}, {idx} segments, {total_duration/3600:.2f}h")

    with open(output_dir / "manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)

    config = {
        "dac": {"hop_length": dac_model.hop_length, "latent_dim": dac_model.latent_dim,
                "sample_rate": 24000},
        "segmentation": SEG_CONFIG,
        "total_segments": idx,
        "total_duration_hours": round(total_duration / 3600, 2),
    }
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nDone! {idx} segments, {total_duration/3600:.2f} hours")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="/home/ubuntu/modi_dataset")
    parser.add_argument("--output", default="/home/ubuntu/lewm-tts/v2/processed_data")
    args = parser.parse_args()
    preprocess_dataset(args.input, args.output)


"""
Preprocess Modi dataset → DAC latents.
Encodes all audio through frozen DAC encoder, saves continuous latents.
"""

import os
import json
import numpy as np
import torch
import torchaudio
import dac
from datasets import load_from_disk
from pathlib import Path
import argparse

SEG_CONFIG = {
    "min_duration": 3.0,
    "max_duration": 12.0,
    "target_duration": 8.0,
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,
}

def detect_silence_points(waveform, sr, threshold_db=-40, min_silence_len=0.3):
    frame_length = int(0.025 * sr)
    hop = int(0.010 * sr)
    n_frames = 1 + (len(waveform) - frame_length) // hop
    rms_db = []
    for i in range(n_frames):
        start = i * hop
        frame = waveform[start:start + frame_length]
        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
        rms_db.append(20 * np.log10(rms + 1e-10))
    rms_db = np.array(rms_db)
    is_silence = rms_db < threshold_db
    silence_points = []
    min_frames = int(min_silence_len / 0.010)
    in_silence = False
    silence_start = 0
    for i, s in enumerate(is_silence):
        if s and not in_silence:
            silence_start = i
            in_silence = True
        elif not s and in_silence:
            if i - silence_start >= min_frames:
                mid_frame = (silence_start + i) // 2
                silence_points.append(mid_frame * hop)
            in_silence = False
    return silence_points

def segment_audio(waveform, sr, text, config):
    total_duration = len(waveform) / sr
    if total_duration <= config["max_duration"]:
        return [(waveform, text)] if total_duration >= config["min_duration"] else []
    silence_points = detect_silence_points(waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"])
    split_points = [0] + silence_points + [len(waveform)]
    segments = []
    current_start = 0
    for point in split_points[1:]:
        current_duration = (point - current_start) / sr
        if current_duration >= config["target_duration"]:
            chunk = waveform[current_start:point]
            chunk_duration = len(chunk) / sr
            if config["min_duration"] <= chunk_duration <= config["max_duration"] * 1.2:
                segments.append(chunk)
            elif chunk_duration > config["max_duration"] * 1.2:
                max_samples = int(config["max_duration"] * sr)
                for j in range(0, len(chunk), max_samples):
                    sub = chunk[j:j + max_samples]
                    if len(sub) / sr >= config["min_duration"]:
                        segments.append(sub)
            current_start = point
    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined
    results = []
    total_samples = sum(len(s) for s in segments)
    char_idx = 0
    for seg in segments:
        proportion = len(seg) / total_samples
        n_chars = max(1, int(len(text) * proportion))
        seg_text = text[char_idx:char_idx + n_chars].strip()
        char_idx += n_chars
        results.append((seg, seg_text))
    return results

def trim_silence(waveform, sr, threshold_db=-45):
    threshold = 10 ** (threshold_db / 20)
    abs_wav = np.abs(waveform)
    window = int(0.01 * sr)
    start = 0
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break
    end = len(abs_wav)
    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break
    pad = int(0.05 * sr)
    return waveform[max(0, start - pad):min(len(waveform), end + pad)]

def preprocess_dataset(input_path, output_path):
    print(f"Loading dataset from {input_path}...")
    ds = load_from_disk(input_path)

    # Load DAC model
    print("Loading DAC 24kHz model...")
    model_path = dac.utils.download(model_type='24khz')
    dac_model = dac.DAC.load(model_path)
    dac_model.eval()
    dac_model = dac_model.to('cuda')
    print(f"DAC loaded. Hop={dac_model.hop_length}, latent_dim={dac_model.latent_dim}")

    output_dir = Path(output_path)
    latent_dir = output_dir / "latents"
    audio_dir = output_dir / "audio"
    latent_dir.mkdir(parents=True, exist_ok=True)
    audio_dir.mkdir(parents=True, exist_ok=True)

    manifest = []
    total_duration = 0.0
    idx = 0

    for i in range(len(ds)):
        sample = ds[i]
        audio_data = sample["audio"]
        waveform = np.array(audio_data["array"], dtype=np.float32)
        sr = audio_data["sampling_rate"]
        text = sample["text"]

        waveform = trim_silence(waveform, sr)
        segments = segment_audio(waveform, sr, text, SEG_CONFIG)

        for seg_wav, seg_text in segments:
            if not seg_text:
                continue
            duration = len(seg_wav) / sr
            total_duration += duration

            # Save audio
            audio_path = audio_dir / f"{idx:06d}.wav"
            wav_tensor = torch.from_numpy(seg_wav).unsqueeze(0)
            torchaudio.save(str(audio_path), wav_tensor, sr)

            # Encode with DAC → continuous latents
            with torch.no_grad():
                wav_input = wav_tensor.unsqueeze(0).to('cuda')  # [1, 1, samples]
                wav_input = dac_model.preprocess(wav_input, sr)
                z = dac_model.encoder(wav_input)  # [1, 1024, T]
                z = z.squeeze(0).cpu()  # [1024, T]

            latent_path = latent_dir / f"{idx:06d}.pt"
            torch.save(z, str(latent_path))

            manifest.append({
                "id": idx,
                "audio_path": str(audio_path),
                "latent_path": str(latent_path),
                "text": seg_text,
                "duration": round(duration, 3),
                "latent_frames": z.shape[1],
            })
            idx += 1

        if (i + 1) % 100 == 0:
            print(f"  Processed {i+1}/{len(ds)}, {idx} segments, {total_duration/3600:.2f}h")

    with open(output_dir / "manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)

    config = {
        "dac": {"hop_length": dac_model.hop_length, "latent_dim": dac_model.latent_dim,
                "sample_rate": 24000},
        "segmentation": SEG_CONFIG,
        "total_segments": idx,
        "total_duration_hours": round(total_duration / 3600, 2),
    }
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nDone! {idx} segments, {total_duration/3600:.2f} hours")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="/home/ubuntu/modi_dataset")
    parser.add_argument("--output", default="/home/ubuntu/lewm-tts/v2/processed_data")
    args = parser.parse_args()
    preprocess_dataset(args.input, args.output)


"""Dataset for DAC-based LeWM TTS."""

import json
import torch
from torch.utils.data import Dataset, DataLoader

class TTSDataset(Dataset):
    def __init__(self, manifest_path, max_latent_frames=900):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)
        self.manifest = [e for e in self.manifest if e["latent_frames"] <= max_latent_frames]
        print(f"Dataset: {len(self.manifest)} samples")

    def __len__(self):
        return len(self.manifest)

    def __getitem__(self, idx):
        entry = self.manifest[idx]
        latent = torch.load(entry["latent_path"], weights_only=True)  # [1024, T]
        latent = latent.transpose(0, 1)  # [T, 1024]
        tokens = list(entry["text"].encode("utf-8"))
        text_tokens = torch.tensor(tokens, dtype=torch.long)
        return {
            "latent": latent,
            "text_tokens": text_tokens,
            "latent_frames": latent.shape[0],
            "text_len": len(tokens),
        }

def collate_fn(batch):
    max_lat = max(s["latent_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    dac_dim = batch[0]["latent"].shape[1]
    B = len(batch)

    lat_padded = torch.zeros(B, max_lat, dac_dim)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    lat_mask = torch.ones(B, max_lat, dtype=torch.bool)
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_lat = s["latent_frames"]
        t_text = s["text_len"]
        lat_padded[i, :t_lat] = s["latent"]
        text_padded[i, :t_text] = s["text_tokens"]
        lat_mask[i, :t_lat] = False
        text_mask[i, :t_text] = False

    return {
        "latent": lat_padded,
        "text_tokens": text_padded,
        "latent_mask": lat_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    dataset = TTSDataset(manifest_path)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle,
        num_workers=num_workers, collate_fn=collate_fn,
        pin_memory=True, drop_last=True,
    )
    return loader, dataset


File created successfully at: /home/ubuntu/lewm-tts/v2/dataset.py

"""Dataset for DAC-based LeWM TTS."""

import json
import torch
from torch.utils.data import Dataset, DataLoader

class TTSDataset(Dataset):
    def __init__(self, manifest_path, max_latent_frames=900):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)
        self.manifest = [e for e in self.manifest if e["latent_frames"] <= max_latent_frames]
        print(f"Dataset: {len(self.manifest)} samples")

    def __len__(self):
        return len(self.manifest)

    def __getitem__(self, idx):
        entry = self.manifest[idx]
        latent = torch.load(entry["latent_path"], weights_only=True)  # [1024, T]
        latent = latent.transpose(0, 1)  # [T, 1024]
        tokens = list(entry["text"].encode("utf-8"))
        text_tokens = torch.tensor(tokens, dtype=torch.long)
        return {
            "latent": latent,
            "text_tokens": text_tokens,
            "latent_frames": latent.shape[0],
            "text_len": len(tokens),
        }

def collate_fn(batch):
    max_lat = max(s["latent_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    dac_dim = batch[0]["latent"].shape[1]
    B = len(batch)

    lat_padded = torch.zeros(B, max_lat, dac_dim)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    lat_mask = torch.ones(B, max_lat, dtype=torch.bool)
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_lat = s["latent_frames"]
        t_text = s["text_len"]
        lat_padded[i, :t_lat] = s["latent"]
        text_padded[i, :t_text] = s["text_tokens"]
        lat_mask[i, :t_lat] = False
        text_mask[i, :t_text] = False

    return {
        "latent": lat_padded,
        "text_tokens": text_padded,
        "latent_mask": lat_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    dataset = TTSDataset(manifest_path)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle,
        num_workers=num_workers, collate_fn=collate_fn,
        pin_memory=True, drop_last=True,
    )
    return loader, dataset