| 1 | """ |
| 2 | Usage: |
| 3 | python prepare_csv_wavs.py /path/to/metadata.csv /output/dataset/path [--pretrain] [--workers N] |
| 4 | |
| 5 | CSV format (header required, "|" delimiter): |
| 6 | audio_file|text |
| 7 | /path/to/wavs/audio_0001.wav|Yo! Hello? Hello? |
| 8 | /path/to/wavs/audio_0002.wav|Hi, how are you doing today? I want to go shopping and buy me some lemons. |
| 9 | |
| 10 | Notes: |
| 11 | - audio_file must be an absolute path. |
| 12 | """ |
| 13 | |
| 14 | import concurrent.futures |
| 15 | import multiprocessing |
| 16 | import os |
| 17 | import shutil |
| 18 | import signal |
| 19 | import subprocess |
| 20 | import sys |
| 21 | from contextlib import contextmanager |
| 22 | |
| 23 | |
| 24 | sys.path.append(os.getcwd()) |
| 25 | |
| 26 | import argparse |
| 27 | import csv |
| 28 | import json |
| 29 | from importlib.resources import files |
| 30 | from pathlib import Path |
| 31 | |
| 32 | import soundfile as sf |
| 33 | import torchaudio |
| 34 | from datasets.arrow_writer import ArrowWriter |
| 35 | from tqdm import tqdm |
| 36 | |
| 37 | from f5_tts.model.utils import convert_char_to_pinyin |
| 38 | |
| 39 | |
| 40 | PRETRAINED_VOCAB_PATH = files("f5_tts").joinpath("../../data/Emilia_ZH_EN_pinyin/vocab.txt") |
| 41 | |
| 42 | # Configuration constants |
| 43 | BATCH_SIZE = 100 # Batch size for text conversion |
| 44 | MAX_WORKERS = max(1, multiprocessing.cpu_count() - 1) # Leave one CPU free |
| 45 | THREAD_NAME_PREFIX = "AudioProcessor" |
| 46 | CHUNK_SIZE = 100 # Number of files to process per worker batch |
| 47 | executor = None # Global executor for cleanup |
| 48 | |
| 49 | |
| 50 | def is_csv_wavs_format(input_path): |
| 51 | fpath = Path(input_path).expanduser() |
| 52 | return fpath.is_file() and fpath.suffix.lower() == ".csv" |
| 53 | |
| 54 | |
| 55 | @contextmanager |
| 56 | def graceful_exit(): |
| 57 | """Context manager for graceful shutdown on signals""" |
| 58 | |
| 59 | def signal_handler(signum, frame): |
| 60 | print("\nReceived signal to terminate. Cleaning up...") |
| 61 | if executor is not None: |
| 62 | print("Shutting down executor...") |
| 63 | executor.shutdown(wait=False, cancel_futures=True) |
| 64 | sys.exit(1) |
| 65 | |
| 66 | # Set up signal handlers |
| 67 | signal.signal(signal.SIGINT, signal_handler) |
| 68 | signal.signal(signal.SIGTERM, signal_handler) |
| 69 | |
| 70 | try: |
| 71 | yield |
| 72 | finally: |
| 73 | if executor is not None: |
| 74 | executor.shutdown(wait=False) |
| 75 | |
| 76 | |
| 77 | def process_audio_file(audio_path, text, polyphone): |
| 78 | """Process a single audio file by checking its existence and extracting duration.""" |
| 79 | if not Path(audio_path).exists(): |
| 80 | print(f"audio {audio_path} not found, skipping") |
| 81 | return None |
| 82 | try: |
| 83 | audio_duration = get_audio_duration(audio_path) |
| 84 | if audio_duration <= 0: |
| 85 | raise ValueError(f"Duration {audio_duration} is non-positive.") |
| 86 | return (audio_path, text, audio_duration) |
| 87 | except Exception as e: |
| 88 | print(f"Warning: Failed to process {audio_path} due to error: {e}. Skipping corrupt file.") |
| 89 | return None |
| 90 | |
| 91 | |
| 92 | def batch_convert_texts(texts, polyphone, batch_size=BATCH_SIZE): |
| 93 | """Convert a list of texts to pinyin in batches.""" |
| 94 | converted_texts = [] |
| 95 | for i in tqdm( |
| 96 | range(0, len(texts), batch_size), |
| 97 | total=(len(texts) + batch_size - 1) // batch_size, |
| 98 | desc="Converting texts to pinyin", |
| 99 | ): |
| 100 | batch = texts[i : i + batch_size] |
| 101 | converted_batch = convert_char_to_pinyin(batch, polyphone=polyphone) |
| 102 | converted_texts.extend(converted_batch) |
| 103 | return converted_texts |
| 104 | |
| 105 | |
| 106 | def prepare_csv_wavs_dir(input_path, num_workers=None): |
| 107 | global executor |
| 108 | if not is_csv_wavs_format(input_path): |
| 109 | raise ValueError(f"input must be a .csv file: {input_path}") |
| 110 | audio_path_text_pairs = read_audio_text_pairs(Path(input_path).expanduser().as_posix()) |
| 111 | |
| 112 | polyphone = True |
| 113 | total_files = len(audio_path_text_pairs) |
| 114 | if total_files == 0: |
| 115 | raise RuntimeError("No valid rows found in CSV.") |
| 116 | |
| 117 | # Use provided worker count or calculate optimal number |
| 118 | worker_count = num_workers if num_workers is not None else min(MAX_WORKERS, total_files) |
| 119 | print(f"\nProcessing {total_files} audio files using {worker_count} workers...") |
| 120 | |
| 121 | with graceful_exit(): |
| 122 | # Initialize thread pool with optimized settings |
| 123 | with concurrent.futures.ThreadPoolExecutor( |
| 124 | max_workers=worker_count, thread_name_prefix=THREAD_NAME_PREFIX |
| 125 | ) as exec: |
| 126 | executor = exec |
| 127 | results = [] |
| 128 | |
| 129 | # Process files in chunks for better efficiency |
| 130 | for i in range(0, len(audio_path_text_pairs), CHUNK_SIZE): |
| 131 | chunk = audio_path_text_pairs[i : i + CHUNK_SIZE] |
| 132 | # Submit futures in order |
| 133 | chunk_futures = [executor.submit(process_audio_file, pair[0], pair[1], polyphone) for pair in chunk] |
| 134 | |
| 135 | # Iterate over futures in the original submission order to preserve ordering |
| 136 | for future in tqdm( |
| 137 | chunk_futures, |
| 138 | total=len(chunk), |
| 139 | desc=f"Processing chunk {i // CHUNK_SIZE + 1}/{(total_files + CHUNK_SIZE - 1) // CHUNK_SIZE}", |
| 140 | ): |
| 141 | try: |
| 142 | result = future.result() |
| 143 | if result is not None: |
| 144 | results.append(result) |
| 145 | except Exception as e: |
| 146 | print(f"Error processing file: {e}") |
| 147 | |
| 148 | executor = None |
| 149 | |
| 150 | # Filter out failed results |
| 151 | processed = [res for res in results if res is not None] |
| 152 | if not processed: |
| 153 | raise RuntimeError("No valid audio files were processed!") |
| 154 | |
| 155 | # Batch process text conversion |
| 156 | raw_texts = [item[1] for item in processed] |
| 157 | converted_texts = batch_convert_texts(raw_texts, polyphone, batch_size=BATCH_SIZE) |
| 158 | |
| 159 | # Prepare final results |
| 160 | sub_result = [] |
| 161 | durations = [] |
| 162 | vocab_set = set() |
| 163 | |
| 164 | for (audio_path, _, duration), conv_text in zip(processed, converted_texts): |
| 165 | sub_result.append({"audio_path": audio_path, "text": conv_text, "duration": duration}) |
| 166 | durations.append(duration) |
| 167 | vocab_set.update(list(conv_text)) |
| 168 | |
| 169 | return sub_result, durations, vocab_set |
| 170 | |
| 171 | |
| 172 | def get_audio_duration(audio_path, timeout=5): |
| 173 | """Get the duration of an audio file in seconds with fallbacks.""" |
| 174 | try: |
| 175 | return sf.info(audio_path).duration |
| 176 | except Exception as e: |
| 177 | print(f"Warning: soundfile failed for {audio_path} with error: {e}. Falling back to ffprobe.") |
| 178 | |
| 179 | try: |
| 180 | cmd = [ |
| 181 | "ffprobe", |
| 182 | "-v", |
| 183 | "error", |
| 184 | "-show_entries", |
| 185 | "format=duration", |
| 186 | "-of", |
| 187 | "default=noprint_wrappers=1:nokey=1", |
| 188 | audio_path, |
| 189 | ] |
| 190 | result = subprocess.run( |
| 191 | cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, timeout=timeout |
| 192 | ) |
| 193 | duration_str = result.stdout.strip() |
| 194 | if duration_str: |
| 195 | return float(duration_str) |
| 196 | raise ValueError("Empty duration string from ffprobe.") |
| 197 | except (subprocess.TimeoutExpired, subprocess.SubprocessError, ValueError) as e: |
| 198 | print(f"Warning: ffprobe failed for {audio_path} with error: {e}. Falling back to torchaudio.info.") |
| 199 | |
| 200 | try: |
| 201 | info = torchaudio.info(audio_path) |
| 202 | if info.sample_rate > 0: |
| 203 | return info.num_frames / info.sample_rate |
| 204 | raise ValueError("Invalid sample_rate from torchaudio.info.") |
| 205 | except Exception as e: |
| 206 | raise RuntimeError(f"failed to get duration for {audio_path}: {e}") |
| 207 | |
| 208 | |
| 209 | def read_audio_text_pairs(csv_file_path): |
| 210 | audio_text_pairs = [] |
| 211 | |
| 212 | csv_path = Path(csv_file_path).expanduser().absolute() |
| 213 | with open(csv_path.as_posix(), mode="r", newline="", encoding="utf-8-sig") as csvfile: |
| 214 | reader = csv.reader(csvfile, delimiter="|") |
| 215 | header = next(reader, None) |
| 216 | if header is None: |
| 217 | return audio_text_pairs |
| 218 | if len(header) < 2 or header[0].strip() != "audio_file" or header[1].strip() != "text": |
| 219 | raise ValueError("CSV header must be: audio_file|text") |
| 220 | for row_idx, row in enumerate(reader, start=2): |
| 221 | if len(row) < 2: |
| 222 | continue |
| 223 | audio_file = row[0].strip() |
| 224 | text = row[1].strip() |
| 225 | if not audio_file: |
| 226 | continue |
| 227 | audio_path = Path(audio_file).expanduser() |
| 228 | if not audio_path.is_absolute(): |
| 229 | raise ValueError(f"audio_file must be an absolute path (row {row_idx}): {audio_file}") |
| 230 | audio_text_pairs.append((audio_path.as_posix(), text)) |
| 231 | |
| 232 | return audio_text_pairs |
| 233 | |
| 234 | |
| 235 | def save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_finetune): |
| 236 | out_dir = Path(out_dir) |
| 237 | out_dir.mkdir(exist_ok=True, parents=True) |
| 238 | print(f"\nSaving to {out_dir} ...") |
| 239 | |
| 240 | raw_arrow_path = out_dir / "raw.arrow" |
| 241 | with ArrowWriter(path=raw_arrow_path.as_posix()) as writer: |
| 242 | for line in tqdm(result, desc="Writing to raw.arrow ..."): |
| 243 | writer.write(line) |
| 244 | writer.finalize() |
| 245 | |
| 246 | # Save durations to JSON |
| 247 | dur_json_path = out_dir / "duration.json" |
| 248 | with open(dur_json_path.as_posix(), "w", encoding="utf-8") as f: |
| 249 | json.dump({"duration": duration_list}, f, ensure_ascii=False) |
| 250 | |
| 251 | # Handle vocab file - write only once based on finetune flag |
| 252 | voca_out_path = out_dir / "vocab.txt" |
| 253 | if is_finetune: |
| 254 | file_vocab_finetune = PRETRAINED_VOCAB_PATH.as_posix() |
| 255 | shutil.copy2(file_vocab_finetune, voca_out_path) |
| 256 | else: |
| 257 | with open(voca_out_path.as_posix(), "w") as f: |
| 258 | for vocab in sorted(text_vocab_set): |
| 259 | f.write(vocab + "\n") |
| 260 | |
| 261 | dataset_name = out_dir.stem |
| 262 | print(f"\nFor {dataset_name}, sample count: {len(result)}") |
| 263 | print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") |
| 264 | print(f"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours") |
| 265 | |
| 266 | |
| 267 | def prepare_and_save_set(inp_dir, out_dir, is_finetune: bool = True, num_workers: int = None): |
| 268 | if is_finetune: |
| 269 | assert PRETRAINED_VOCAB_PATH.exists(), f"pretrained vocab.txt not found: {PRETRAINED_VOCAB_PATH}" |
| 270 | sub_result, durations, vocab_set = prepare_csv_wavs_dir(inp_dir, num_workers=num_workers) |
| 271 | save_prepped_dataset(out_dir, sub_result, durations, vocab_set, is_finetune) |
| 272 | |
| 273 | |
| 274 | def get_args(): |
| 275 | parser = argparse.ArgumentParser(description="Prepare and save dataset.") |
| 276 | parser.add_argument( |
| 277 | "inp_dir", |
| 278 | type=str, |
| 279 | help="Input CSV with header 'audio_file|text' and absolute wav paths.", |
| 280 | ) |
| 281 | parser.add_argument("out_dir", type=str, help="Output directory to save the prepared data.") |
| 282 | parser.add_argument("--pretrain", action="store_true", help="Enable for new pretrain, otherwise is a fine-tune") |
| 283 | parser.add_argument("--workers", type=int, help=f"Number of worker threads (default: {MAX_WORKERS})") |
| 284 | return parser.parse_args() |
| 285 | |
| 286 | |
| 287 | def cli(): |
| 288 | try: |
| 289 | args = get_args() |
| 290 | prepare_and_save_set(args.inp_dir, args.out_dir, is_finetune=not args.pretrain, num_workers=args.workers) |
| 291 | except KeyboardInterrupt: |
| 292 | print("\nOperation cancelled by user. Cleaning up...") |
| 293 | if executor is not None: |
| 294 | executor.shutdown(wait=False, cancel_futures=True) |
| 295 | sys.exit(1) |
| 296 | |
| 297 | |
| 298 | if __name__ == "__main__": |
| 299 | cli() |
| 300 |