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)