Spaces:
Runtime error
Runtime error
| import os | |
| import glob | |
| import torch | |
| import torchaudio | |
| from tqdm import tqdm | |
| from src.chatterbox_.tts_turbo import ChatterboxTurboTTS | |
| from src.chatterbox_.tts import ChatterboxTTS, punc_norm | |
| from src.chatterbox_.models.s3tokenizer import S3_SR | |
| from src.utils import setup_logger | |
| from src.config import TrainConfig | |
| logger = setup_logger(__name__) | |
| def preprocess_dataset_file_based(config, tts_engine: ChatterboxTTS): | |
| """ | |
| Reads .wav and .txt file pairs in a folder, processes them, and saves them as .pt. | |
| Structure: | |
| ID.wav (Audio) | |
| ID.txt (Text) | |
| """ | |
| 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) | |
| tts_engine.ve.eval() | |
| tts_engine.s3gen.eval() | |
| search_path = os.path.join(config.wav_dir, "*.wav") | |
| wav_files = glob.glob(search_path) | |
| if len(wav_files) == 0: | |
| logger.error(f"ERROR: No .wav files found in folder '{config.wav_dir}'!") | |
| return | |
| logger.info(f"Processing dataset... Found audio file: {len(wav_files)}") | |
| success_count = 0 | |
| SPEECH_STOP_ID = getattr(tts_engine.t3.hp, 'stop_speech_token', 6562) | |
| for wav_path in tqdm(wav_files, desc="Preprocessing"): | |
| try: | |
| filename = os.path.basename(wav_path) | |
| file_id = os.path.splitext(filename)[0] | |
| txt_path = os.path.join(config.wav_dir, f"{file_id}.txt") | |
| if not os.path.exists(txt_path): | |
| logger.warning(f"Text file not found, skipping: {file_id}") | |
| continue | |
| with open(txt_path, "r", encoding="utf-8") as f: | |
| raw_text = f.read().strip() | |
| if not raw_text: | |
| continue | |
| wav, sr = torchaudio.load(wav_path) | |
| if wav.shape[0] > 1: | |
| wav = wav.mean(dim=0, keepdim=True) | |
| if sr != S3_SR: | |
| resampler = torchaudio.transforms.Resample(sr, S3_SR) | |
| wav = resampler(wav) | |
| wav = wav.to(device) | |
| with torch.no_grad(): | |
| wav_np = wav.cpu().squeeze().numpy() | |
| spk_emb_np = tts_engine.ve.embeds_from_wavs([wav_np], sample_rate=S3_SR) | |
| speaker_emb = torch.from_numpy(spk_emb_np[0]).cpu() | |
| s_tokens, _ = tts_engine.s3gen.tokenizer(wav.unsqueeze(0)) | |
| raw_speech_tokens = s_tokens.squeeze().cpu() | |
| stop_speech_tensor = torch.tensor([SPEECH_STOP_ID], dtype=raw_speech_tokens.dtype) | |
| speech_tokens = torch.cat([raw_speech_tokens, stop_speech_tensor], dim=0) | |
| prompt_samples = int(config.prompt_duration * S3_SR) | |
| if wav.shape[1] < prompt_samples: | |
| prompt_wav = torch.nn.functional.pad(wav, (0, prompt_samples - wav.shape[1])) | |
| else: | |
| prompt_wav = wav[:, :prompt_samples] | |
| p_tokens, _ = tts_engine.s3gen.tokenizer(prompt_wav.unsqueeze(0)) | |
| prompt_tokens = p_tokens.squeeze().cpu() | |
| clean_text = punc_norm(raw_text) | |
| if config.is_turbo: | |
| token_output = tts_engine.tokenizer(clean_text, return_tensors="pt") | |
| raw_text_tokens = token_output.input_ids[0].cpu() | |
| if tts_engine.tokenizer.eos_token_id is not None: | |
| text_eos = torch.tensor([tts_engine.tokenizer.eos_token_id], dtype=raw_text_tokens.dtype) | |
| text_tokens = torch.cat([raw_text_tokens, text_eos], dim=0) | |
| else: | |
| text_tokens = raw_text_tokens | |
| else: | |
| text_tokens = tts_engine.tokenizer.text_to_tokens(clean_text).squeeze(0).cpu() | |
| # --- 5. SAVING --- | |
| # We keep the file name: ID.pt | |
| save_path = os.path.join(config.preprocessed_dir, f"{file_id}.pt") | |
| torch.save({ | |
| "speech_tokens": speech_tokens, | |
| "speaker_emb": speaker_emb, | |
| "prompt_tokens": prompt_tokens, | |
| "text_tokens": text_tokens, | |
| }, save_path) | |
| success_count += 1 | |
| except Exception as e: | |
| logger.error(f"Error ({filename}): {e}") | |
| continue | |
| logger.info(f"Preprocessing completed! Success: {success_count}/{len(wav_files)}") | |
| 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_file_based(cfg, tts_engine) |