| 1 | # Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation) |
| 2 | |
| 3 | import argparse |
| 4 | import ast |
| 5 | import json |
| 6 | import os |
| 7 | import sys |
| 8 | |
| 9 | |
| 10 | sys.path.append(os.getcwd()) |
| 11 | |
| 12 | import multiprocessing as mp |
| 13 | from importlib.resources import files |
| 14 | |
| 15 | import numpy as np |
| 16 | |
| 17 | from f5_tts.eval.utils_eval import get_librispeech_test, run_asr_wer, run_sim |
| 18 | |
| 19 | |
| 20 | rel_path = str(files("f5_tts").joinpath("../../")) |
| 21 | |
| 22 | |
| 23 | def get_args(): |
| 24 | parser = argparse.ArgumentParser() |
| 25 | parser.add_argument("-e", "--eval_task", type=str, default="wer", choices=["sim", "wer"]) |
| 26 | parser.add_argument("-l", "--lang", type=str, default="en") |
| 27 | parser.add_argument("-g", "--gen_wav_dir", type=str, required=True) |
| 28 | parser.add_argument("-p", "--librispeech_test_clean_path", type=str, required=True) |
| 29 | parser.add_argument( |
| 30 | "-n", "--gpu_nums", type=str, default="8", help="Number of GPUs to use (e.g., 8) or GPU list (e.g., [0,1,2,3])" |
| 31 | ) |
| 32 | parser.add_argument("--local", action="store_true", help="Use local custom checkpoint directory") |
| 33 | return parser.parse_args() |
| 34 | |
| 35 | |
| 36 | def parse_gpu_nums(gpu_nums_str): |
| 37 | try: |
| 38 | if gpu_nums_str.startswith("[") and gpu_nums_str.endswith("]"): |
| 39 | gpu_list = ast.literal_eval(gpu_nums_str) |
| 40 | if isinstance(gpu_list, list): |
| 41 | return gpu_list |
| 42 | return list(range(int(gpu_nums_str))) |
| 43 | except (ValueError, SyntaxError): |
| 44 | raise argparse.ArgumentTypeError( |
| 45 | f"Invalid GPU specification: {gpu_nums_str}. Use a number (e.g., 8) or a list (e.g., [0,1,2,3])" |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | def main(): |
| 50 | args = get_args() |
| 51 | eval_task = args.eval_task |
| 52 | lang = args.lang |
| 53 | librispeech_test_clean_path = args.librispeech_test_clean_path # test-clean path |
| 54 | gen_wav_dir = args.gen_wav_dir |
| 55 | metalst = rel_path + "/data/librispeech_pc_test_clean_cross_sentence.lst" |
| 56 | |
| 57 | gpus = parse_gpu_nums(args.gpu_nums) |
| 58 | test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path) |
| 59 | |
| 60 | ## In LibriSpeech, some speakers utilized varying voice characteristics for different characters in the book, |
| 61 | ## leading to a low similarity for the ground truth in some cases. |
| 62 | # test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = True) # eval ground truth |
| 63 | |
| 64 | local = args.local |
| 65 | if local: # use local custom checkpoint dir |
| 66 | asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3" |
| 67 | else: |
| 68 | asr_ckpt_dir = "" # auto download to cache dir |
| 69 | wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth" |
| 70 | |
| 71 | # -------------------------------------------------------------------------- |
| 72 | |
| 73 | full_results = [] |
| 74 | metrics = [] |
| 75 | |
| 76 | if eval_task == "wer": |
| 77 | with mp.Pool(processes=len(gpus)) as pool: |
| 78 | args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set] |
| 79 | results = pool.map(run_asr_wer, args) |
| 80 | for r in results: |
| 81 | full_results.extend(r) |
| 82 | elif eval_task == "sim": |
| 83 | with mp.Pool(processes=len(gpus)) as pool: |
| 84 | args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set] |
| 85 | results = pool.map(run_sim, args) |
| 86 | for r in results: |
| 87 | full_results.extend(r) |
| 88 | else: |
| 89 | raise ValueError(f"Unknown metric type: {eval_task}") |
| 90 | |
| 91 | result_path = f"{gen_wav_dir}/_{eval_task}_results.jsonl" |
| 92 | with open(result_path, "w") as f: |
| 93 | for line in full_results: |
| 94 | metrics.append(line[eval_task]) |
| 95 | f.write(json.dumps(line, ensure_ascii=False) + "\n") |
| 96 | metric = round(np.mean(metrics), 5) |
| 97 | f.write(f"\n{eval_task.upper()}: {metric}\n") |
| 98 | |
| 99 | print(f"\nTotal {len(metrics)} samples") |
| 100 | print(f"{eval_task.upper()}: {metric}") |
| 101 | print(f"{eval_task.upper()} results saved to {result_path}") |
| 102 | |
| 103 | |
| 104 | if __name__ == "__main__": |
| 105 | main() |
| 106 |