| 1 | import os |
| 2 | import sys |
| 3 | |
| 4 | |
| 5 | sys.path.append(os.getcwd()) |
| 6 | |
| 7 | import json |
| 8 | from concurrent.futures import ProcessPoolExecutor |
| 9 | from importlib.resources import files |
| 10 | from pathlib import Path |
| 11 | |
| 12 | import soundfile as sf |
| 13 | from datasets.arrow_writer import ArrowWriter |
| 14 | from tqdm import tqdm |
| 15 | |
| 16 | |
| 17 | def deal_with_audio_dir(audio_dir): |
| 18 | sub_result, durations = [], [] |
| 19 | vocab_set = set() |
| 20 | audio_lists = list(audio_dir.rglob("*.wav")) |
| 21 | |
| 22 | for line in audio_lists: |
| 23 | text_path = line.with_suffix(".normalized.txt") |
| 24 | text = open(text_path, "r").read().strip() |
| 25 | duration = sf.info(line).duration |
| 26 | if duration < 0.4 or duration > 30: |
| 27 | continue |
| 28 | sub_result.append({"audio_path": str(line), "text": text, "duration": duration}) |
| 29 | durations.append(duration) |
| 30 | vocab_set.update(list(text)) |
| 31 | return sub_result, durations, vocab_set |
| 32 | |
| 33 | |
| 34 | def main(): |
| 35 | result = [] |
| 36 | duration_list = [] |
| 37 | text_vocab_set = set() |
| 38 | |
| 39 | # process raw data |
| 40 | executor = ProcessPoolExecutor(max_workers=max_workers) |
| 41 | futures = [] |
| 42 | |
| 43 | for subset in tqdm(SUB_SET): |
| 44 | dataset_path = Path(os.path.join(dataset_dir, subset)) |
| 45 | [ |
| 46 | futures.append(executor.submit(deal_with_audio_dir, audio_dir)) |
| 47 | for audio_dir in dataset_path.iterdir() |
| 48 | if audio_dir.is_dir() |
| 49 | ] |
| 50 | for future in tqdm(futures, total=len(futures)): |
| 51 | sub_result, durations, vocab_set = future.result() |
| 52 | result.extend(sub_result) |
| 53 | duration_list.extend(durations) |
| 54 | text_vocab_set.update(vocab_set) |
| 55 | executor.shutdown() |
| 56 | |
| 57 | # save preprocessed dataset to disk |
| 58 | if not os.path.exists(f"{save_dir}"): |
| 59 | os.makedirs(f"{save_dir}") |
| 60 | print(f"\nSaving to {save_dir} ...") |
| 61 | |
| 62 | with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer: |
| 63 | for line in tqdm(result, desc="Writing to raw.arrow ..."): |
| 64 | writer.write(line) |
| 65 | writer.finalize() |
| 66 | |
| 67 | # dup a json separately saving duration in case for DynamicBatchSampler ease |
| 68 | with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f: |
| 69 | json.dump({"duration": duration_list}, f, ensure_ascii=False) |
| 70 | |
| 71 | # vocab map, i.e. tokenizer |
| 72 | with open(f"{save_dir}/vocab.txt", "w") as f: |
| 73 | for vocab in sorted(text_vocab_set): |
| 74 | f.write(vocab + "\n") |
| 75 | |
| 76 | print(f"\nFor {dataset_name}, sample count: {len(result)}") |
| 77 | print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") |
| 78 | print(f"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours") |
| 79 | |
| 80 | |
| 81 | if __name__ == "__main__": |
| 82 | max_workers = 36 |
| 83 | |
| 84 | tokenizer = "char" # "pinyin" | "char" |
| 85 | |
| 86 | SUB_SET = ["train-clean-100", "train-clean-360", "train-other-500"] |
| 87 | dataset_dir = "<SOME_PATH>/LibriTTS" |
| 88 | dataset_name = f"LibriTTS_{'_'.join(SUB_SET)}_{tokenizer}".replace("train-clean-", "").replace("train-other-", "") |
| 89 | save_dir = str(files("f5_tts").joinpath("../../")) + f"/data/{dataset_name}" |
| 90 | print(f"\nPrepare for {dataset_name}, will save to {save_dir}\n") |
| 91 | main() |
| 92 | |
| 93 | # For LibriTTS_100_360_500_char, sample count: 354218 |
| 94 | # For LibriTTS_100_360_500_char, vocab size is: 78 |
| 95 | # For LibriTTS_100_360_500_char, total 554.09 hours |
| 96 |