chatterbox-finetuning/src/preprocess_ljspeech.py

import os
import queue
import threading
import torch
import torchaudio
import numpy as np
import pandas as pd
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed
from src.chatterbox_.tts import ChatterboxTTS, punc_norm
from src.chatterbox_.tts_turbo import ChatterboxTurboTTS
from src.chatterbox_.models.s3tokenizer import S3_SR
from src.utils import setup_logger
from src.config import TrainConfig

logger = setup_logger(__name__)

# ── tunables (override via env vars) ─────────────────────────────────────────
LOAD_WORKERS    = int(os.environ.get("PP_LOAD_WORKERS", min(16, os.cpu_count() or 8)))
SAVE_WORKERS    = int(os.environ.get("PP_SAVE_WORKERS", 4))
BATCH_SIZE      = int(os.environ.get("PP_BATCH_SIZE",   32))
PREFETCH_BATCHES = int(os.environ.get("PP_PREFETCH",    3))   # batches to load ahead
# ─────────────────────────────────────────────────────────────────────────────


def _load_wav(wav_path: str, target_sr: int):
    """Load + resample a WAV. Runs in a thread pool (I/O bound)."""
    try:
        wav, sr = torchaudio.load(wav_path)
        if wav.shape[0] > 1:
            wav = wav.mean(dim=0, keepdim=True)
        if sr != target_sr:
            wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
        return wav   # (1, T) float32 CPU
    except Exception as e:
        logger.error(f"Load error {wav_path}: {e}")
        return None


def _save_pt(data: dict, path: str):
    torch.save(data, path)


def _tokenize_text(args):
    """Tokenize one text sample. Runs in a thread (CPU, releases GIL via C ext)."""
    row, is_turbo, tokenizer, punc_norm_fn = args
    raw_text  = str(row[2]) if len(row) > 2 else str(row[1])
    clean_txt = punc_norm_fn(raw_text)
    if is_turbo:
        tok_out = tokenizer(clean_txt, return_tensors="pt")
        toks = tok_out.input_ids[0].cpu()
        if tokenizer.eos_token_id is not None:
            eos  = torch.tensor([tokenizer.eos_token_id], dtype=toks.dtype)
            toks = torch.cat([toks, eos])
    else:
        toks = tokenizer.text_to_tokens(clean_txt).squeeze(0).cpu()
    return toks


def preprocess_dataset_ljspeech(config, tts_engine):

    data = pd.read_csv(config.csv_path, sep="|", header=None, quoting=3)
    os.makedirs(config.preprocessed_dir, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tts_engine.ve.to(device)
    tts_engine.s3gen.to(device)

    SPEECH_STOP_ID = getattr(tts_engine.t3.hp, 'stop_speech_token', 6562)
    prompt_samples = int(config.prompt_duration * S3_SR)

    # ── skip already-done .pt files ───────────────────────────────────────────
    rows = []
    for _, row in data.iterrows():
        filename = str(row[0])
        if not filename.endswith(".wav"):
            filename += ".wav"
        save_path = os.path.join(
            config.preprocessed_dir, filename.replace(".wav", ".pt")
        )
        if not os.path.exists(save_path):
            rows.append((row, filename, save_path))

    skipped = len(data) - len(rows)
    logger.info(f"Total: {len(data)} | Skipping: {skipped} | To process: {len(rows)}")
    if not rows:
        logger.info("Nothing to do.")
        return

    # ── prefetch thread: loads WAV batches into a queue while GPU runs ────────
    # This gives true CPU/GPU overlap: GPU never waits on disk.
    prefetch_q = queue.Queue(maxsize=PREFETCH_BATCHES)

    def _prefetch_worker():
        with ThreadPoolExecutor(max_workers=LOAD_WORKERS) as loader:
            for b in range(0, len(rows), BATCH_SIZE):
                batch = rows[b : b + BATCH_SIZE]
                wav_futs = {
                    loader.submit(
                        _load_wav,
                        os.path.join(config.wav_dir, item[1]),
                        S3_SR,
                    ): item
                    for item in batch
                }
                loaded = []
                for fut in as_completed(wav_futs):
                    wav = fut.result()
                    if wav is not None:
                        loaded.append((wav_futs[fut], wav))
                prefetch_q.put(loaded)
        prefetch_q.put(None)  # sentinel

    prefetch_thread = threading.Thread(target=_prefetch_worker, daemon=True)
    prefetch_thread.start()

    # ── text tokenizer pool: runs in parallel with GPU ops ────────────────────
    text_pool   = ThreadPoolExecutor(max_workers=LOAD_WORKERS)
    saver_pool  = ThreadPoolExecutor(max_workers=SAVE_WORKERS)
    pending_saves = []
    success_count = skipped

    with tqdm(total=len(rows), desc=f"Preprocessing (batch={BATCH_SIZE})") as pbar:
        while True:
            loaded = prefetch_q.get()
            if loaded is None:
                break
            if not loaded:
                continue

            items, wavs = zip(*loaded)
            rows_b = [it[0] for it in items]
            saves_b = [it[2] for it in items]
            B = len(wavs)

            # ── 1. Batch speaker embeddings (GPU) ─────────────────────────────
            wav_nps = [w.squeeze().numpy() for w in wavs]
            with torch.no_grad():
                spk_embs_np = tts_engine.ve.embeds_from_wavs(
                    wav_nps, sample_rate=S3_SR
                )

            # ── 2a. Batch prompt tokens (fixed length → always safe to batch) ─
            prompt_batch = torch.zeros(B, 1, prompt_samples)
            for i, w in enumerate(wavs):
                if w.shape[1] < prompt_samples:
                    prompt_batch[i, :, : w.shape[1]] = w
                else:
                    prompt_batch[i] = w[:, :prompt_samples]
            prompt_batch = prompt_batch.to(device)
            with torch.no_grad():
                p_tokens_batch, _ = tts_engine.s3gen.tokenizer(prompt_batch)
            p_tokens_batch = p_tokens_batch.cpu()   # (B, L_prompt)

            # ── 2b. Batch speech tokens (variable length — pad + trim) ─────────
            max_len = max(w.shape[1] for w in wavs)
            speech_batch = torch.zeros(B, 1, max_len)
            for i, w in enumerate(wavs):
                speech_batch[i, :, : w.shape[1]] = w
            speech_batch = speech_batch.to(device)
            with torch.no_grad():
                s_tokens_batch, _ = tts_engine.s3gen.tokenizer(speech_batch)
            s_tokens_batch = s_tokens_batch.cpu()   # (B, L_speech)

            # ── 3. Text tokens (parallel CPU threads, runs during GPU above) ──
            # NOTE: submitted BEFORE GPU steps so threads run concurrently
            text_futs = [
                text_pool.submit(
                    _tokenize_text,
                    (row, config.is_turbo, tts_engine.tokenizer, punc_norm),
                )
                for row in rows_b
            ]

            # ── 4. Assemble + save each sample ────────────────────────────────
            stop = torch.tensor([SPEECH_STOP_ID], dtype=torch.long)
            for i in range(B):
                try:
                    # Trim speech tokens to actual audio length.
                    # The tokenizer downsamples at a fixed rate; tokens beyond
                    # the real audio correspond to padding silence — drop them.
                    actual_frames = wavs[i].shape[1]
                    total_frames  = max_len
                    total_toks    = s_tokens_batch.shape[1]
                    keep = max(1, round(total_toks * actual_frames / total_frames))
                    raw_speech    = s_tokens_batch[i, :keep]
                    speech_tokens = torch.cat([raw_speech, stop])
                    prompt_tokens = p_tokens_batch[i]
                    speaker_emb   = torch.from_numpy(spk_embs_np[i]).cpu()
                    text_tokens   = text_futs[i].result()

                    _save_pt(
                        {
                            "speech_tokens": speech_tokens,
                            "speaker_emb":   speaker_emb,
                            "prompt_tokens": prompt_tokens,
                            "text_tokens":   text_tokens,
                        },
                        saves_b[i],
                    )
                    success_count += 1

                except Exception as e:
                    logger.error(f"Error (sample {i} in batch): {e}")

                pbar.update(1)

            # Flush completed saves to bound memory
            done_saves = [f for f in pending_saves if f.done()]
            for f in done_saves:
                try:
                    f.result()
                except Exception as e:
                    logger.error(f"Save error: {e}")
            pending_saves = [f for f in pending_saves if not f.done()]

    # Wait for any remaining background saves
    for f in as_completed(pending_saves):
        try:
            f.result()
        except Exception as e:
            logger.error(f"Save error: {e}")

    prefetch_thread.join(timeout=5)
    text_pool.shutdown(wait=False)
    saver_pool.shutdown(wait=True)

    logger.info(f"Preprocessing done. Success: {success_count}/{len(data)}")


if __name__ == "__main__":

    cfg = TrainConfig()

    if cfg.is_turbo:
        EngineClass = ChatterboxTurboTTS
    else:
        EngineClass = ChatterboxTTS

    logger.info(f"{EngineClass} engine starting...")
    tts_engine = EngineClass.from_local(cfg.model_dir, device="cpu")

    preprocess_dataset_ljspeech(cfg, tts_engine)