     1→"""
     2→Data preprocessing for LeWM TTS.
     3→- Loads Modi dataset from HuggingFace Arrow format
     4→- Segments long clips into 3-15 second chunks using silence detection
     5→- Extracts mel spectrograms
     6→- Saves processed dataset
     7→"""
     8→
     9→import os
    10→import json
    11→import numpy as np
    12→import torch
    13→import torchaudio
    14→from datasets import load_from_disk
    15→from pathlib import Path
    16→import argparse
    17→
    18→
    19→# ─── Mel spectrogram config ───
    20→MEL_CONFIG = {
    21→    "sample_rate": 24000,
    22→    "n_fft": 1024,
    23→    "hop_length": 256,
    24→    "n_mels": 80,
    25→    "f_min": 0,
    26→    "f_max": 12000,
    27→    "power": 1.0,
    28→}
    29→
    30→# ─── Segmentation config ───
    31→SEG_CONFIG = {
    32→    "min_duration": 3.0,   # seconds
    33→    "max_duration": 12.0,  # seconds
    34→    "target_duration": 8.0,  # preferred chunk size
    35→    "silence_threshold_db": -40,
    36→    "min_silence_len": 0.3,  # seconds
    37→}
    38→
    39→
    40→def detect_silence_points(waveform: np.ndarray, sr: int,
    41→                          threshold_db: float = -40,
    42→                          min_silence_len: float = 0.3) -> list:
    43→    """Find silence points in audio for splitting."""
    44→    # Convert to dB
    45→    frame_length = int(0.025 * sr)  # 25ms frames
    46→    hop = int(0.010 * sr)           # 10ms hop
    47→
    48→    # RMS energy per frame
    49→    n_frames = 1 + (len(waveform) - frame_length) // hop
    50→    rms_db = []
    51→    for i in range(n_frames):
    52→        start = i * hop
    53→        frame = waveform[start:start + frame_length]
    54→        rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
    55→        rms_db.append(20 * np.log10(rms + 1e-10))
    56→
    57→    rms_db = np.array(rms_db)
    58→    is_silence = rms_db < threshold_db
    59→
    60→    # Find contiguous silence regions
    61→    silence_points = []
    62→    min_frames = int(min_silence_len / 0.010)
    63→
    64→    in_silence = False
    65→    silence_start = 0
    66→    for i, s in enumerate(is_silence):
    67→        if s and not in_silence:
    68→            silence_start = i
    69→            in_silence = True
    70→        elif not s and in_silence:
    71→            if i - silence_start >= min_frames:
    72→                # Midpoint of silence region in samples
    73→                mid_frame = (silence_start + i) // 2
    74→                mid_sample = mid_frame * hop
    75→                silence_points.append(mid_sample)
    76→            in_silence = False
    77→
    78→    return silence_points
    79→
    80→
    81→def segment_audio(waveform: np.ndarray, sr: int, text: str, config: dict) -> list:
    82→    """Segment a long audio clip into shorter chunks at silence boundaries."""
    83→    total_duration = len(waveform) / sr
    84→
    85→    # If already short enough, return as-is
    86→    if total_duration <= config["max_duration"]:
    87→        if total_duration >= config["min_duration"]:
    88→            return [(waveform, text)]
    89→        else:
    90→            return []  # Too short
    91→
    92→    # Find silence points for natural splitting
    93→    silence_points = detect_silence_points(
    94→        waveform, sr,
    95→        threshold_db=config["silence_threshold_db"],
    96→        min_silence_len=config["min_silence_len"]
    97→    )
    98→
    99→    # Add start and end
   100→    split_points = [0] + silence_points + [len(waveform)]
   101→
   102→    # Greedily merge segments to hit target duration
   103→    segments = []
   104→    current_start = 0
   105→
   106→    for i, point in enumerate(split_points[1:], 1):
   107→        current_duration = (point - current_start) / sr
   108→
   109→        if current_duration >= config["target_duration"]:
   110→            # Take this segment
   111→            chunk = waveform[current_start:point]
   112→            chunk_duration = len(chunk) / sr
   113→            if config["min_duration"] <= chunk_duration <= config["max_duration"] * 1.2:
   114→                segments.append(chunk)
   115→            elif chunk_duration > config["max_duration"] * 1.2:
   116→                # Force split at max_duration
   117→                max_samples = int(config["max_duration"] * sr)
   118→                for j in range(0, len(chunk), max_samples):
   119→                    sub = chunk[j:j + max_samples]
   120→                    if len(sub) / sr >= config["min_duration"]:
   121→                        segments.append(sub)
   122→            current_start = point
   123→
   124→    # Handle remaining audio
   125→    if current_start < len(waveform):
   126→        remaining = waveform[current_start:]
   127→        if len(remaining) / sr >= config["min_duration"]:
   128→            segments.append(remaining)
   129→        elif segments:
   130→            # Append to last segment if it won't make it too long
   131→            combined = np.concatenate([segments[-1], remaining])
   132→            if len(combined) / sr <= config["max_duration"] * 1.2:
   133→                segments[-1] = combined
   134→
   135→    # For text: we can't perfectly align without forced alignment,
   136→    # so we use the full text for all segments (will be refined later with forced alignment)
   137→    # For now, distribute text roughly proportionally
   138→    results = []
   139→    total_samples = sum(len(s) for s in segments)
   140→    text_chars = list(text)
   141→    char_idx = 0
   142→
   143→    for seg in segments:
   144→        proportion = len(seg) / total_samples
   145→        n_chars = max(1, int(len(text) * proportion))
   146→        seg_text = text[char_idx:char_idx + n_chars]
   147→        char_idx += n_chars
   148→        results.append((seg, seg_text.strip()))
   149→
   150→    return results
   151→
   152→
   153→def extract_mel(waveform: torch.Tensor, config: dict) -> torch.Tensor:
   154→    """Extract log-mel spectrogram."""
   155→    mel_transform = torchaudio.transforms.MelSpectrogram(
   156→        sample_rate=config["sample_rate"],
   157→        n_fft=config["n_fft"],
   158→        hop_length=config["hop_length"],
   159→        n_mels=config["n_mels"],
   160→        f_min=config["f_min"],
   161→        f_max=config["f_max"],
   162→        power=config["power"],
   163→    )
   164→
   165→    mel = mel_transform(waveform)
   166→    # Log scale
   167→    log_mel = torch.log(torch.clamp(mel, min=1e-5))
   168→    return log_mel
   169→
   170→
   171→def trim_silence(waveform: np.ndarray, sr: int, threshold_db: float = -45) -> np.ndarray:
   172→    """Trim leading and trailing silence."""
   173→    # Convert to amplitude threshold
   174→    threshold = 10 ** (threshold_db / 20)
   175→
   176→    # Find non-silent regions
   177→    abs_wav = np.abs(waveform)
   178→
   179→    # Use a sliding window
   180→    window = int(0.01 * sr)  # 10ms
   181→    for start in range(0, len(abs_wav) - window, window):
   182→        if np.max(abs_wav[start:start + window]) > threshold:
   183→            break
   184→
   185→    for end in range(len(abs_wav), window, -window):
   186→        if np.max(abs_wav[end - window:end]) > threshold:
   187→            break
   188→
   189→    # Add small padding
   190→    pad = int(0.05 * sr)  # 50ms padding
   191→    start = max(0, start - pad)
   192→    end = min(len(waveform), end + pad)
   193→
   194→    return waveform[start:end]
   195→
   196→
   197→def preprocess_dataset(input_path: str, output_path: str):
   198→    """Main preprocessing pipeline."""
   199→    print(f"Loading dataset from {input_path}...")
   200→    ds = load_from_disk(input_path)
   201→
   202→    output_dir = Path(output_path)
   203→    audio_dir = output_dir / "audio"
   204→    mel_dir = output_dir / "mels"
   205→    audio_dir.mkdir(parents=True, exist_ok=True)
   206→    mel_dir.mkdir(parents=True, exist_ok=True)
   207→
   208→    manifest = []
   209→    total_duration = 0.0
   210→    idx = 0
   211→
   212→    for i in range(len(ds)):
   213→        sample = ds[i]
   214→        audio_data = sample["audio"]
   215→        waveform = np.array(audio_data["array"], dtype=np.float32)
   216→        sr = audio_data["sampling_rate"]
   217→        text = sample["text"]
   218→
   219→        # Trim silence
   220→        waveform = trim_silence(waveform, sr)
   221→
   222→        # Segment into chunks
   223→        segments = segment_audio(waveform, sr, text, SEG_CONFIG)
   224→
   225→        for seg_wav, seg_text in segments:
   226→            if not seg_text:
   227→                continue
   228→
   229→            duration = len(seg_wav) / sr
   230→            total_duration += duration
   231→
   232→            # Save audio
   233→            audio_path = audio_dir / f"{idx:06d}.wav"
   234→            wav_tensor = torch.from_numpy(seg_wav).unsqueeze(0)
   235→            torchaudio.save(str(audio_path), wav_tensor, sr)
   236→
   237→            # Extract and save mel
   238→            mel = extract_mel(wav_tensor, MEL_CONFIG)
   239→            mel_path = mel_dir / f"{idx:06d}.pt"
   240→            torch.save(mel.squeeze(0), str(mel_path))  # [n_mels, T]
   241→
   242→            manifest.append({
   243→                "id": idx,
   244→                "audio_path": str(audio_path),
   245→                "mel_path": str(mel_path),
   246→                "text": seg_text,
   247→                "duration": round(duration, 3),
   248→                "mel_frames": mel.shape[-1],
   249→            })
   250→
   251→            idx += 1
   252→
   253→        if (i + 1) % 100 == 0:
   254→            print(f"  Processed {i+1}/{len(ds)} samples, {idx} segments, {total_duration/3600:.2f}h total")
   255→
   256→    # Save manifest
   257→    manifest_path = output_dir / "manifest.json"
   258→    with open(manifest_path, "w", encoding="utf-8") as f:
   259→        json.dump(manifest, f, ensure_ascii=False, indent=2)
   260→
   261→    # Save config
   262→    config = {
   263→        "mel": MEL_CONFIG,
   264→        "segmentation": SEG_CONFIG,
   265→        "total_segments": idx,
   266→        "total_duration_hours": round(total_duration / 3600, 2),
   267→        "original_samples": len(ds),
   268→    }
   269→    config_path = output_dir / "config.json"
   270→    with open(config_path, "w") as f:
   271→        json.dump(config, f, indent=2)
   272→
   273→    print(f"\nDone! {idx} segments, {total_duration/3600:.2f} hours")
   274→    print(f"Manifest: {manifest_path}")
   275→    print(f"Config: {config_path}")
   276→
   277→
   278→if __name__ == "__main__":
   279→    parser = argparse.ArgumentParser()
   280→    parser.add_argument("--input", default="/home/ubuntu/modi_dataset")
   281→    parser.add_argument("--output", default="/home/ubuntu/lewm-tts/processed_data")
   282→    args = parser.parse_args()
   283→    preprocess_dataset(args.input, args.output)
   284→

"""
Data preprocessing for LeWM TTS.
- Loads Modi dataset from HuggingFace Arrow format
- Segments long clips into 3-15 second chunks using silence detection
- Extracts mel spectrograms
- Saves processed dataset
"""

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 spectrogram config ───
MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 80,
    "f_min": 0,
    "f_max": 12000,
    "power": 1.0,
}

# ─── Segmentation config ───
SEG_CONFIG = {
    "min_duration": 3.0,   # seconds
    "max_duration": 12.0,  # seconds
    "target_duration": 8.0,  # preferred chunk size
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,  # seconds
}

def detect_silence_points(waveform: np.ndarray, sr: int,
                          threshold_db: float = -40,
                          min_silence_len: float = 0.3) -> list:
    """Find silence points in audio for splitting."""
    # Convert to dB
    frame_length = int(0.025 * sr)  # 25ms frames
    hop = int(0.010 * sr)           # 10ms hop

    # RMS energy per frame
    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

    # Find contiguous silence regions
    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:
                # Midpoint of silence region in samples
                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: np.ndarray, sr: int, text: str, config: dict) -> list:
    """Segment a long audio clip into shorter chunks at silence boundaries."""
    total_duration = len(waveform) / sr

    # If already short enough, return as-is
    if total_duration <= config["max_duration"]:
        if total_duration >= config["min_duration"]:
            return [(waveform, text)]
        else:
            return []  # Too short

    # Find silence points for natural splitting
    silence_points = detect_silence_points(
        waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"]
    )

    # Add start and end
    split_points = [0] + silence_points + [len(waveform)]

    # Greedily merge segments to hit target duration
    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"]:
            # Take this segment
            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:
                # Force split at max_duration
                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

    # Handle remaining audio
    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            # Append to last segment if it won't make it too long
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined

    # For text: we can't perfectly align without forced alignment,
    # so we use the full text for all segments (will be refined later with forced alignment)
    # For now, distribute text roughly proportionally
    results = []
    total_samples = sum(len(s) for s in segments)
    text_chars = list(text)
    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: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram."""
    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 scale
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def trim_silence(waveform: np.ndarray, sr: int, threshold_db: float = -45) -> np.ndarray:
    """Trim leading and trailing silence."""
    # Convert to amplitude threshold
    threshold = 10 ** (threshold_db / 20)

    # Find non-silent regions
    abs_wav = np.abs(waveform)

    # Use a sliding window
    window = int(0.01 * sr)  # 10ms
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break

    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break

    # Add small padding
    pad = int(0.05 * sr)  # 50ms padding
    start = max(0, start - pad)
    end = min(len(waveform), end + pad)

    return waveform[start:end]

def preprocess_dataset(input_path: str, output_path: str):
    """Main preprocessing pipeline."""
    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"]

        # Trim silence
        waveform = trim_silence(waveform, sr)

        # Segment into chunks
        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)

            # Extract and save mel
            mel = extract_mel(wav_tensor, MEL_CONFIG)
            mel_path = mel_dir / f"{idx:06d}.pt"
            torch.save(mel.squeeze(0), str(mel_path))  # [n_mels, T]

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

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

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

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

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


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

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

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

from model import build_model
from dataset import build_dataloader

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            optimizer.step()
            scheduler.step()

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

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

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

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

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

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

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

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

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


     1→"""
     2→Dataset and DataLoader for LeWM TTS training.
     3→"""
     4→
     5→import json
     6→import torch
     7→from torch.utils.data import Dataset, DataLoader
     8→from pathlib import Path
     9→
    10→
    11→class TTSDataset(Dataset):
    12→    """Loads preprocessed mel spectrograms and text for TTS training."""
    13→
    14→    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
    15→        with open(manifest_path, "r", encoding="utf-8") as f:
    16→            self.manifest = json.load(f)
    17→
    18→        self.max_mel_frames = max_mel_frames
    19→        self.max_text_len = max_text_len
    20→
    21→        # Filter out entries that are too long
    22→        self.manifest = [
    23→            e for e in self.manifest
    24→            if e["mel_frames"] <= max_mel_frames
    25→        ]
    26→        print(f"Dataset: {len(self.manifest)} samples loaded")
    27→
    28→    def __len__(self):
    29→        return len(self.manifest)
    30→
    31→    def text_to_tokens(self, text):
    32→        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
    33→        # UTF-8 byte encoding — simple, universal, no tokenizer needed
    34→        tokens = list(text.encode("utf-8"))
    35→        return tokens
    36→
    37→    def __getitem__(self, idx):
    38→        entry = self.manifest[idx]
    39→
    40→        # Load mel
    41→        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]
    42→
    43→        # Tokenize text
    44→        tokens = self.text_to_tokens(entry["text"])
    45→        tokens = tokens[:self.max_text_len]
    46→        text_tokens = torch.tensor(tokens, dtype=torch.long)
    47→
    48→        return {
    49→            "mel": mel,
    50→            "text_tokens": text_tokens,
    51→            "mel_frames": mel.shape[1],
    52→            "text_len": len(tokens),
    53→        }
    54→
    55→
    56→def collate_fn(batch):
    57→    """Pad and batch samples."""
    58→    # Find max lengths
    59→    max_mel = max(s["mel_frames"] for s in batch)
    60→    max_text = max(s["text_len"] for s in batch)
    61→    n_mels = batch[0]["mel"].shape[0]
    62→    B = len(batch)
    63→
    64→    # Allocate padded tensors
    65→    mel_padded = torch.zeros(B, n_mels, max_mel)
    66→    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    67→    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
    68→    text_mask = torch.ones(B, max_text, dtype=torch.bool)
    69→
    70→    for i, s in enumerate(batch):
    71→        t_mel = s["mel_frames"]
    72→        t_text = s["text_len"]
    73→        mel_padded[i, :, :t_mel] = s["mel"]
    74→        text_padded[i, :t_text] = s["text_tokens"]
    75→        mel_mask[i, :t_mel] = False
    76→        text_mask[i, :t_text] = False
    77→
    78→    return {
    79→        "mel": mel_padded,
    80→        "text_tokens": text_padded,
    81→        "mel_mask": mel_mask,
    82→        "text_mask": text_mask,
    83→    }
    84→
    85→
    86→def build_dataloader(manifest_path, batch_size=16, num_workers=4, shuffle=True):
    87→    dataset = TTSDataset(manifest_path)
    88→    loader = DataLoader(
    89→        dataset,
    90→        batch_size=batch_size,
    91→        shuffle=shuffle,
    92→        num_workers=num_workers,
    93→        collate_fn=collate_fn,
    94→        pin_memory=True,
    95→        drop_last=True,
    96→    )
    97→    return loader, dataset
    98→

"""
Dataset and DataLoader for LeWM TTS training.
"""

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

class TTSDataset(Dataset):
    """Loads preprocessed mel spectrograms and text for TTS training."""

    def __init__(self, manifest_path, max_mel_frames=1200, max_text_len=512):
        with open(manifest_path, "r", encoding="utf-8") as f:
            self.manifest = json.load(f)

        self.max_mel_frames = max_mel_frames
        self.max_text_len = max_text_len

        # Filter out entries that are too long
        self.manifest = [
            e for e in self.manifest
            if e["mel_frames"] <= max_mel_frames
        ]
        print(f"Dataset: {len(self.manifest)} samples loaded")

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

    def text_to_tokens(self, text):
        """Convert text to byte-level tokens (handles any Unicode including Devanagari)."""
        # UTF-8 byte encoding — simple, universal, no tokenizer needed
        tokens = list(text.encode("utf-8"))
        return tokens

    def __getitem__(self, idx):
        entry = self.manifest[idx]

        # Load mel
        mel = torch.load(entry["mel_path"], weights_only=True)  # [n_mels, T]

        # Tokenize text
        tokens = self.text_to_tokens(entry["text"])
        tokens = tokens[:self.max_text_len]
        text_tokens = torch.tensor(tokens, dtype=torch.long)

        return {
            "mel": mel,
            "text_tokens": text_tokens,
            "mel_frames": mel.shape[1],
            "text_len": len(tokens),
        }

def collate_fn(batch):
    """Pad and batch samples."""
    # Find max lengths
    max_mel = max(s["mel_frames"] for s in batch)
    max_text = max(s["text_len"] for s in batch)
    n_mels = batch[0]["mel"].shape[0]
    B = len(batch)

    # Allocate padded tensors
    mel_padded = torch.zeros(B, n_mels, max_mel)
    text_padded = torch.zeros(B, max_text, dtype=torch.long)
    mel_mask = torch.ones(B, max_mel, dtype=torch.bool)   # True = padding
    text_mask = torch.ones(B, max_text, dtype=torch.bool)

    for i, s in enumerate(batch):
        t_mel = s["mel_frames"]
        t_text = s["text_len"]
        mel_padded[i, :, :t_mel] = s["mel"]
        text_padded[i, :t_text] = s["text_tokens"]
        mel_mask[i, :t_mel] = False
        text_mask[i, :t_text] = False

    return {
        "mel": mel_padded,
        "text_tokens": text_padded,
        "mel_mask": mel_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


ToolSearch

Question: select:TaskCreate


Question: select:TaskCreate


Question: select:TaskUpdate


Question: select:TaskUpdate


TaskCreate

Change preprocess.py: n_mels 80 → 100, match Vocos config

Update MEL_CONFIG in preprocess.py to n_mels=100, f_max=None (let it auto), power=2.0 to match Vocos mel-24khz extraction params


Update model.py: n_mels default 80 → 100, add speaker embedding

Change all n_mels=80 defaults to 100 in AudioEncoder, MelDecoder, build_model. Add speaker embedding support for multi-speaker.


Update train.py: n_mels config to 100, add speaker_id support

Change hardcoded n_mels=80 to 100 in training config


Update inference.py: use Vocos directly (no Griffin-Lim)

Remove GL vocoder, use Vocos.decode() directly since model now outputs 100 mels matching Vocos


Update inference_engine.py: Vocos direct decode, 100 mels

Remove GL from Vocoder class, use Vocos.decode() directly. Update mel extraction for prompts to 100 mels.


Reprocess dataset with 100 mels

Run preprocess.py to regenerate all mel spectrograms with n_mels=100. Save to processed_data_v3 to not overwrite old data.


Verify: load new mels, run model forward, Vocos decode end-to-end

Sanity check that the full pipeline works: load 100-mel → model forward → mel decoder output → Vocos decode → clean audio


Task #1 created successfully: Update inference_engine.py: Vocos direct decode, 100 mels


Task #2 created successfully: Update inference.py: use Vocos directly (no Griffin-Lim)


Task #3 created successfully: Update model.py: n_mels default 80 → 100, add speaker embedding


Task #4 created successfully: Verify: load new mels, run model forward, Vocos decode end-to-end


Task #5 created successfully: Update train.py: n_mels config to 100, add speaker_id support


Task #6 created successfully: Change preprocess.py: n_mels 80 → 100, match Vocos config


Task #7 created successfully: Reprocess dataset with 100 mels

Reprocess dataset with 100 mels


First, let me check the exact Vocos mel config so we match it perfectly:


Updated task #6 status


Bash

python3 -c "
from vocos import Vocos
v = Vocos.from_pretrained('charactr/vocos-mel-24khz')
fe = v.feature_extractor
ms = fe.mel_spec
print('=== Vocos mel-24khz feature extractor config ===')
print(f'sample_rate: {ms.sample_rate}')
print(f'n_fft: {ms.n_fft}')
print(f'win_length: {ms.win_length}')
print(f'hop_length: {ms.hop_length}')
print(f'n_mels: {ms.n_mels}')
print(f'f_min: {ms.f_min}')
print(f'f_max: {ms.f_max}')
print(f'power: {ms.power}')
print(f'normalized: {ms.normalized}')
print(f'pad: {fe.padding}')

# Check how Vocos computes features
import inspect
print()
print('=== feature_extractor.forward ===')
print(inspect.getsource(fe.forward))
" 2>&1 | grep -v Warning

Get exact Vocos mel spectrogram config


=== Vocos mel-24khz feature extractor config ===
sample_rate: 24000
n_fft: 1024
win_length: 1024
hop_length: 256
n_mels: 100
f_min: 0.0
f_max: None
power: 1
normalized: False
pad: center

=== feature_extractor.forward ===
    def forward(self, audio, **kwargs):
        if self.padding == "same":
            pad = self.mel_spec.win_length - self.mel_spec.hop_length
            audio = torch.nn.functional.pad(audio, (pad // 2, pad // 2), mode="reflect")
        mel = self.mel_spec(audio)
        features = safe_log(mel)
        return features


python3 -c "
from vocos.feature_extractors import safe_log
import inspect
print(inspect.getsource(safe_log))
"

Check safe_log implementation


def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
    """
    Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.

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

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


Now I know the exact Vocos config:
- `n_mels=100, n_fft=1024, hop_length=256, sr=24000, f_min=0, f_max=None, power=1`
- Features = `log(clamp(mel, min=1e-7))` — basically same as our `log(clamp(mel, min=1e-5))`

Let me update everything. Starting with preprocess.py:


Edit

MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 100,
    "f_min": 0,
    "f_max": None,      # full bandwidth (matches Vocos mel-24khz)
    "power": 1.0,
    "clip_val": 1e-7,   # log clamp value (matches Vocos safe_log)
}


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

MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 100,
    "f_min": 0,
    "f_max": None,      # full bandwidth (matches Vocos mel-24khz)
    "power": 1.0,
    "clip_val": 1e-7,   # log clamp value (matches Vocos safe_log)
}

"""
Data preprocessing for LeWM TTS.
- Loads Modi dataset from HuggingFace Arrow format
- Segments long clips into 3-15 second chunks using silence detection
- Extracts mel spectrograms
- Saves processed dataset
"""

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 spectrogram config ───
MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 80,
    "f_min": 0,
    "f_max": 12000,
    "power": 1.0,
}

# ─── Segmentation config ───
SEG_CONFIG = {
    "min_duration": 3.0,   # seconds
    "max_duration": 12.0,  # seconds
    "target_duration": 8.0,  # preferred chunk size
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,  # seconds
}

def detect_silence_points(waveform: np.ndarray, sr: int,
                          threshold_db: float = -40,
                          min_silence_len: float = 0.3) -> list:
    """Find silence points in audio for splitting."""
    # Convert to dB
    frame_length = int(0.025 * sr)  # 25ms frames
    hop = int(0.010 * sr)           # 10ms hop

    # RMS energy per frame
    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

    # Find contiguous silence regions
    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:
                # Midpoint of silence region in samples
                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: np.ndarray, sr: int, text: str, config: dict) -> list:
    """Segment a long audio clip into shorter chunks at silence boundaries."""
    total_duration = len(waveform) / sr

    # If already short enough, return as-is
    if total_duration <= config["max_duration"]:
        if total_duration >= config["min_duration"]:
            return [(waveform, text)]
        else:
            return []  # Too short

    # Find silence points for natural splitting
    silence_points = detect_silence_points(
        waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"]
    )

    # Add start and end
    split_points = [0] + silence_points + [len(waveform)]

    # Greedily merge segments to hit target duration
    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"]:
            # Take this segment
            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:
                # Force split at max_duration
                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

    # Handle remaining audio
    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            # Append to last segment if it won't make it too long
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined

    # For text: we can't perfectly align without forced alignment,
    # so we use the full text for all segments (will be refined later with forced alignment)
    # For now, distribute text roughly proportionally
    results = []
    total_samples = sum(len(s) for s in segments)
    text_chars = list(text)
    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: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram."""
    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 scale
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def trim_silence(waveform: np.ndarray, sr: int, threshold_db: float = -45) -> np.ndarray:
    """Trim leading and trailing silence."""
    # Convert to amplitude threshold
    threshold = 10 ** (threshold_db / 20)

    # Find non-silent regions
    abs_wav = np.abs(waveform)

    # Use a sliding window
    window = int(0.01 * sr)  # 10ms
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break

    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break

    # Add small padding
    pad = int(0.05 * sr)  # 50ms padding
    start = max(0, start - pad)
    end = min(len(waveform), end + pad)

    return waveform[start:end]

def preprocess_dataset(input_path: str, output_path: str):
    """Main preprocessing pipeline."""
    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"]

        # Trim silence
        waveform = trim_silence(waveform, sr)

        # Segment into chunks
        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)

            # Extract and save mel
            mel = extract_mel(wav_tensor, MEL_CONFIG)
            mel_path = mel_dir / f"{idx:06d}.pt"
            torch.save(mel.squeeze(0), str(mel_path))  # [n_mels, T]

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

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

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

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

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

+    "f_max": None,      # full bandwidth (matches Vocos mel-24khz)

+    "clip_val": 1e-7,   # log clamp value (matches Vocos safe_log)


Update the extract_mel function to match Vocos exactly:


def extract_mel(waveform: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram."""
    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 scale
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def extract_mel(waveform: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram (matches Vocos mel-24khz config)."""
    mel_kwargs = dict(
        sample_rate=config["sample_rate"],
        n_fft=config["n_fft"],
        hop_length=config["hop_length"],
        n_mels=config["n_mels"],
        power=config["power"],
    )
    if config.get("f_min") is not None:
        mel_kwargs["f_min"] = config["f_min"]
    if config.get("f_max") is not None:
        mel_kwargs["f_max"] = config["f_max"]

    mel_transform = torchaudio.transforms.MelSpectrogram(**mel_kwargs)
    mel = mel_transform(waveform)
    clip_val = config.get("clip_val", 1e-7)
    log_mel = torch.log(torch.clamp(mel, min=clip_val))
    return log_mel


def extract_mel(waveform: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram (matches Vocos mel-24khz config)."""
    mel_kwargs = dict(
        sample_rate=config["sample_rate"],
        n_fft=config["n_fft"],
        hop_length=config["hop_length"],
        n_mels=config["n_mels"],
        power=config["power"],
    )
    if config.get("f_min") is not None:
        mel_kwargs["f_min"] = config["f_min"]
    if config.get("f_max") is not None:
        mel_kwargs["f_max"] = config["f_max"]

    mel_transform = torchaudio.transforms.MelSpectrogram(**mel_kwargs)
    mel = mel_transform(waveform)
    clip_val = config.get("clip_val", 1e-7)
    log_mel = torch.log(torch.clamp(mel, min=clip_val))
    return log_mel

"""
Data preprocessing for LeWM TTS.
- Loads Modi dataset from HuggingFace Arrow format
- Segments long clips into 3-15 second chunks using silence detection
- Extracts mel spectrograms
- Saves processed dataset
"""

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 spectrogram config ───
MEL_CONFIG = {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 100,
    "f_min": 0,
    "f_max": None,      # full bandwidth (matches Vocos mel-24khz)
    "power": 1.0,
    "clip_val": 1e-7,   # log clamp value (matches Vocos safe_log)
}

# ─── Segmentation config ───
SEG_CONFIG = {
    "min_duration": 3.0,   # seconds
    "max_duration": 12.0,  # seconds
    "target_duration": 8.0,  # preferred chunk size
    "silence_threshold_db": -40,
    "min_silence_len": 0.3,  # seconds
}

def detect_silence_points(waveform: np.ndarray, sr: int,
                          threshold_db: float = -40,
                          min_silence_len: float = 0.3) -> list:
    """Find silence points in audio for splitting."""
    # Convert to dB
    frame_length = int(0.025 * sr)  # 25ms frames
    hop = int(0.010 * sr)           # 10ms hop

    # RMS energy per frame
    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

    # Find contiguous silence regions
    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:
                # Midpoint of silence region in samples
                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: np.ndarray, sr: int, text: str, config: dict) -> list:
    """Segment a long audio clip into shorter chunks at silence boundaries."""
    total_duration = len(waveform) / sr

    # If already short enough, return as-is
    if total_duration <= config["max_duration"]:
        if total_duration >= config["min_duration"]:
            return [(waveform, text)]
        else:
            return []  # Too short

    # Find silence points for natural splitting
    silence_points = detect_silence_points(
        waveform, sr,
        threshold_db=config["silence_threshold_db"],
        min_silence_len=config["min_silence_len"]
    )

    # Add start and end
    split_points = [0] + silence_points + [len(waveform)]

    # Greedily merge segments to hit target duration
    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"]:
            # Take this segment
            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:
                # Force split at max_duration
                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

    # Handle remaining audio
    if current_start < len(waveform):
        remaining = waveform[current_start:]
        if len(remaining) / sr >= config["min_duration"]:
            segments.append(remaining)
        elif segments:
            # Append to last segment if it won't make it too long
            combined = np.concatenate([segments[-1], remaining])
            if len(combined) / sr <= config["max_duration"] * 1.2:
                segments[-1] = combined

    # For text: we can't perfectly align without forced alignment,
    # so we use the full text for all segments (will be refined later with forced alignment)
    # For now, distribute text roughly proportionally
    results = []
    total_samples = sum(len(s) for s in segments)
    text_chars = list(text)
    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: torch.Tensor, config: dict) -> torch.Tensor:
    """Extract log-mel spectrogram."""
    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 scale
    log_mel = torch.log(torch.clamp(mel, min=1e-5))
    return log_mel

def trim_silence(waveform: np.ndarray, sr: int, threshold_db: float = -45) -> np.ndarray:
    """Trim leading and trailing silence."""
    # Convert to amplitude threshold
    threshold = 10 ** (threshold_db / 20)

    # Find non-silent regions
    abs_wav = np.abs(waveform)

    # Use a sliding window
    window = int(0.01 * sr)  # 10ms
    for start in range(0, len(abs_wav) - window, window):
        if np.max(abs_wav[start:start + window]) > threshold:
            break

    for end in range(len(abs_wav), window, -window):
        if np.max(abs_wav[end - window:end]) > threshold:
            break

    # Add small padding
    pad = int(0.05 * sr)  # 50ms padding
    start = max(0, start - pad)
    end = min(len(waveform), end + pad)

    return waveform[start:end]

def preprocess_dataset(input_path: str, output_path: str):
    """Main preprocessing pipeline."""
    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"]

        # Trim silence
        waveform = trim_silence(waveform, sr)

        # Segment into chunks
        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)

            # Extract and save mel
            mel = extract_mel(wav_tensor, MEL_CONFIG)
            mel_path = mel_dir / f"{idx:06d}.pt"
            torch.save(mel.squeeze(0), str(mel_path))  # [n_mels, T]

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

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

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

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

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

-    mel_transform = torchaudio.transforms.MelSpectrogram(

+    """Extract log-mel spectrogram (matches Vocos mel-24khz config)."""

+    if config.get("f_min") is not None:

+        mel_kwargs["f_min"] = config["f_min"]

+    if config.get("f_max") is not None:

+        mel_kwargs["f_max"] = config["f_max"]

+    mel_transform = torchaudio.transforms.MelSpectrogram(**mel_kwargs)

-    log_mel = torch.log(torch.clamp(mel, min=1e-5))

+    clip_val = config.get("clip_val", 1e-7)

+    log_mel = torch.log(torch.clamp(mel, min=clip_val))


Updated task #3 status


Now update model.py — change all n_mels defaults to 100 and add speaker embedding:


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return z, mu, logvar

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

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

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

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

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

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

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

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

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

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

        predicted = self.output_proj(predicted)
        return predicted

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

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

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

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

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

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

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

        return out, k, v

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

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

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

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

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

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

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

        return out, k, v

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

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

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

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

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

        return self.output_proj(x), cache

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

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

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

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

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

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

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

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

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

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

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

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

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

        self.config = config

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-    def __init__(self, n_mels=80, d_model=256, nhead=4, num_layers=4,

+    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,

-    def __init__(self, d_model=256, n_mels=80):

+    def __init__(self, d_model=256, n_mels=100):


Now add speaker embedding to the model. The cleanest approach: a learned speaker embedding that gets added to the text embeddings (conditions the predictor on speaker identity):


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

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

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

        self.config = config

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

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

        # Mel decoder: reconstructs mel from latent embeddings
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
        self.recon_weight = config.get("recon_weight", 1.0)

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

        self.config = config


<tool_use_error>String to replace not found in file.
String:     def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)

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

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

        self.config = config</tool_use_error>

Error: String to replace not found in file.
String:     def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)

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

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

        self.config = config


   365→    Two losses:
   366→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   367→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   368→    """
   369→
   370→    def __init__(self, config):
   371→        super().__init__()
   372→        d_model = config.get("d_model", 256)
   373→        nhead = config.get("nhead", 4)
   374→        n_mels = config.get("n_mels", 80)
   375→        text_vocab = config.get("text_vocab_size", 256)
   376→        text_layers = config.get("text_encoder_layers", 4)
   377→        audio_layers = config.get("audio_encoder_layers", 4)
   378→        predictor_layers = config.get("predictor_layers", 6)
   379→        dropout = config.get("dropout", 0.1)
   380→        self.kl_weight = config.get("kl_weight", 0.01)
   381→
   382→        self.text_encoder = TextEncoder(
   383→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   384→            num_layers=text_layers, dropout=dropout,
   385→        )
   386→        self.audio_encoder = AudioEncoder(
   387→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   388→            num_layers=audio_layers, dropout=dropout,
   389→        )
   390→        self.predictor = JEPAPredictor(
   391→            d_model=d_model, nhead=nhead,
   392→            num_layers=predictor_layers, dropout=dropout,
   393→        )
   394→
   395→        # Mel decoder: reconstructs mel from latent embeddings
   396→        # Needed so the latent space is grounded to actual audio
   397→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
   398→        self.recon_weight = config.get("recon_weight", 1.0)
   399→
   400→        self.config = config
   401→
   402→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
   403→        """
   404→        Args:

    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

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

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

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

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
        """
        Args:


The `replace_all` for n_mels missed the one inside `config.get("n_mels", 80)`. Let me fix: