prabaerode/zero-shot-tts
0
1import math2import os3import random4import string5 6import torch7import torch.nn.functional as F8import torchaudio9from tqdm import tqdm10 11from f5_tts.eval.ecapa_tdnn import ECAPA_TDNN_SMALL12from f5_tts.model.modules import MelSpec13from f5_tts.model.utils import convert_char_to_pinyin14 15 16# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav17def get_seedtts_testset_metainfo(metalst):18 f = open(metalst)19 lines = f.readlines()20 f.close()21 metainfo = []22 for line in lines:23 if len(line.strip().split("|")) == 5:24 utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|")25 elif len(line.strip().split("|")) == 4:26 utt, prompt_text, prompt_wav, gt_text = line.strip().split("|")27 gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav")28 if not os.path.isabs(prompt_wav):29 prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)30 metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav))31 return metainfo32 33 34# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav35def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path):36 f = open(metalst)37 lines = f.readlines()38 f.close()39 metainfo = []40 for line in lines:41 ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t")42 43 # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc)44 ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-")45 ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac")46 47 # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc)48 gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-")49 gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac")50 51 metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav))52 53 return metainfo54 55 56# padded to max length mel batch57def padded_mel_batch(ref_mels):58 max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax()59 padded_ref_mels = []60 for mel in ref_mels:61 padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value=0)62 padded_ref_mels.append(padded_ref_mel)63 padded_ref_mels = torch.stack(padded_ref_mels)64 padded_ref_mels = padded_ref_mels.permute(0, 2, 1)65 return padded_ref_mels66 67 68# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav69 70 71def get_inference_prompt(72 metainfo,73 speed=1.0,74 tokenizer="pinyin",75 polyphone=True,76 target_sample_rate=24000,77 n_fft=1024,78 win_length=1024,79 n_mel_channels=100,80 hop_length=256,81 mel_spec_type="vocos",82 target_rms=0.1,83 use_truth_duration=False,84 infer_batch_size=1,85 num_buckets=200,86 min_secs=3,87 max_secs=40,88):89 prompts_all = []90 91 min_tokens = min_secs * target_sample_rate // hop_length92 max_tokens = max_secs * target_sample_rate // hop_length93 94 batch_accum = [0] * num_buckets95 utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = (96 [[] for _ in range(num_buckets)] for _ in range(6)97 )98 99 mel_spectrogram = MelSpec(100 n_fft=n_fft,101 hop_length=hop_length,102 win_length=win_length,103 n_mel_channels=n_mel_channels,104 target_sample_rate=target_sample_rate,105 mel_spec_type=mel_spec_type,106 )107 108 for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."):109 # Audio110 ref_audio, ref_sr = torchaudio.load(prompt_wav)111 ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio)))112 if ref_rms < target_rms:113 ref_audio = ref_audio * target_rms / ref_rms114 assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue."115 if ref_sr != target_sample_rate:116 resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)117 ref_audio = resampler(ref_audio)118 119 # Text120 if len(prompt_text[-1].encode("utf-8")) == 1:121 prompt_text = prompt_text + " "122 text = [prompt_text + gt_text]123 if tokenizer == "pinyin":124 text_list = convert_char_to_pinyin(text, polyphone=polyphone)125 else:126 text_list = text127 128 # Duration, mel frame length129 ref_mel_len = ref_audio.shape[-1] // hop_length130 if use_truth_duration:131 gt_audio, gt_sr = torchaudio.load(gt_wav)132 if gt_sr != target_sample_rate:133 resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate)134 gt_audio = resampler(gt_audio)135 total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed)136 137 # # test vocoder resynthesis138 # ref_audio = gt_audio139 else:140 ref_text_len = len(prompt_text.encode("utf-8"))141 gen_text_len = len(gt_text.encode("utf-8"))142 total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)143 144 # to mel spectrogram145 ref_mel = mel_spectrogram(ref_audio)146 ref_mel = ref_mel.squeeze(0)147 148 # deal with batch149 assert infer_batch_size > 0, "infer_batch_size should be greater than 0."150 assert (151 min_tokens <= total_mel_len <= max_tokens152 ), f"Audio {utt} has duration {total_mel_len*hop_length//target_sample_rate}s out of range [{min_secs}, {max_secs}]."153 bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets)154 155 utts[bucket_i].append(utt)156 ref_rms_list[bucket_i].append(ref_rms)157 ref_mels[bucket_i].append(ref_mel)158 ref_mel_lens[bucket_i].append(ref_mel_len)159 total_mel_lens[bucket_i].append(total_mel_len)160 final_text_list[bucket_i].extend(text_list)161 162 batch_accum[bucket_i] += total_mel_len163 164 if batch_accum[bucket_i] >= infer_batch_size:165 # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}")166 prompts_all.append(167 (168 utts[bucket_i],169 ref_rms_list[bucket_i],170 padded_mel_batch(ref_mels[bucket_i]),171 ref_mel_lens[bucket_i],172 total_mel_lens[bucket_i],173 final_text_list[bucket_i],174 )175 )176 batch_accum[bucket_i] = 0177 (178 utts[bucket_i],179 ref_rms_list[bucket_i],180 ref_mels[bucket_i],181 ref_mel_lens[bucket_i],182 total_mel_lens[bucket_i],183 final_text_list[bucket_i],184 ) = [], [], [], [], [], []185 186 # add residual187 for bucket_i, bucket_frames in enumerate(batch_accum):188 if bucket_frames > 0:189 prompts_all.append(190 (191 utts[bucket_i],192 ref_rms_list[bucket_i],193 padded_mel_batch(ref_mels[bucket_i]),194 ref_mel_lens[bucket_i],195 total_mel_lens[bucket_i],196 final_text_list[bucket_i],197 )198 )199 # not only leave easy work for last workers200 random.seed(666)201 random.shuffle(prompts_all)202 203 return prompts_all204 205 206# get wav_res_ref_text of seed-tts test metalst207# https://github.com/BytedanceSpeech/seed-tts-eval208 209 210def get_seed_tts_test(metalst, gen_wav_dir, gpus):211 f = open(metalst)212 lines = f.readlines()213 f.close()214 215 test_set_ = []216 for line in tqdm(lines):217 if len(line.strip().split("|")) == 5:218 utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|")219 elif len(line.strip().split("|")) == 4:220 utt, prompt_text, prompt_wav, gt_text = line.strip().split("|")221 222 if not os.path.exists(os.path.join(gen_wav_dir, utt + ".wav")):223 continue224 gen_wav = os.path.join(gen_wav_dir, utt + ".wav")225 if not os.path.isabs(prompt_wav):226 prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)227 228 test_set_.append((gen_wav, prompt_wav, gt_text))229 230 num_jobs = len(gpus)231 if num_jobs == 1:232 return [(gpus[0], test_set_)]233 234 wav_per_job = len(test_set_) // num_jobs + 1235 test_set = []236 for i in range(num_jobs):237 test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))238 239 return test_set240 241 242# get librispeech test-clean cross sentence test243 244 245def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth=False):246 f = open(metalst)247 lines = f.readlines()248 f.close()249 250 test_set_ = []251 for line in tqdm(lines):252 ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t")253 254 if eval_ground_truth:255 gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-")256 gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac")257 else:258 if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + ".wav")):259 raise FileNotFoundError(f"Generated wav not found: {gen_utt}")260 gen_wav = os.path.join(gen_wav_dir, gen_utt + ".wav")261 262 ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-")263 ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac")264 265 test_set_.append((gen_wav, ref_wav, gen_txt))266 267 num_jobs = len(gpus)268 if num_jobs == 1:269 return [(gpus[0], test_set_)]270 271 wav_per_job = len(test_set_) // num_jobs + 1272 test_set = []273 for i in range(num_jobs):274 test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))275 276 return test_set277 278 279# load asr model280 281 282def load_asr_model(lang, ckpt_dir=""):283 if lang == "zh":284 from funasr import AutoModel285 286 model = AutoModel(287 model=os.path.join(ckpt_dir, "paraformer-zh"),288 # vad_model = os.path.join(ckpt_dir, "fsmn-vad"),289 # punc_model = os.path.join(ckpt_dir, "ct-punc"),290 # spk_model = os.path.join(ckpt_dir, "cam++"),291 disable_update=True,292 ) # following seed-tts setting293 elif lang == "en":294 from faster_whisper import WhisperModel295 296 model_size = "large-v3" if ckpt_dir == "" else ckpt_dir297 model = WhisperModel(model_size, device="cuda", compute_type="float16")298 return model299 300 301# WER Evaluation, the way Seed-TTS does302 303 304def run_asr_wer(args):305 rank, lang, test_set, ckpt_dir = args306 307 if lang == "zh":308 import zhconv309 310 torch.cuda.set_device(rank)311 elif lang == "en":312 os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)313 else:314 raise NotImplementedError(315 "lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now."316 )317 318 asr_model = load_asr_model(lang, ckpt_dir=ckpt_dir)319 320 from zhon.hanzi import punctuation321 322 punctuation_all = punctuation + string.punctuation323 wers = []324 325 from jiwer import compute_measures326 327 for gen_wav, prompt_wav, truth in tqdm(test_set):328 if lang == "zh":329 res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True)330 hypo = res[0]["text"]331 hypo = zhconv.convert(hypo, "zh-cn")332 elif lang == "en":333 segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en")334 hypo = ""335 for segment in segments:336 hypo = hypo + " " + segment.text337 338 # raw_truth = truth339 # raw_hypo = hypo340 341 for x in punctuation_all:342 truth = truth.replace(x, "")343 hypo = hypo.replace(x, "")344 345 truth = truth.replace(" ", " ")346 hypo = hypo.replace(" ", " ")347 348 if lang == "zh":349 truth = " ".join([x for x in truth])350 hypo = " ".join([x for x in hypo])351 elif lang == "en":352 truth = truth.lower()353 hypo = hypo.lower()354 355 measures = compute_measures(truth, hypo)356 wer = measures["wer"]357 358 # ref_list = truth.split(" ")359 # subs = measures["substitutions"] / len(ref_list)360 # dele = measures["deletions"] / len(ref_list)361 # inse = measures["insertions"] / len(ref_list)362 363 wers.append(wer)364 365 return wers366 367 368# SIM Evaluation369 370 371def run_sim(args):372 rank, test_set, ckpt_dir = args373 device = f"cuda:{rank}"374 375 model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type="wavlm_large", config_path=None)376 state_dict = torch.load(ckpt_dir, weights_only=True, map_location=lambda storage, loc: storage)377 model.load_state_dict(state_dict["model"], strict=False)378 379 use_gpu = True if torch.cuda.is_available() else False380 if use_gpu:381 model = model.cuda(device)382 model.eval()383 384 sim_list = []385 for wav1, wav2, truth in tqdm(test_set):386 wav1, sr1 = torchaudio.load(wav1)387 wav2, sr2 = torchaudio.load(wav2)388 389 resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000)390 resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000)391 wav1 = resample1(wav1)392 wav2 = resample2(wav2)393 394 if use_gpu:395 wav1 = wav1.cuda(device)396 wav2 = wav2.cuda(device)397 with torch.no_grad():398 emb1 = model(wav1)399 emb2 = model(wav2)400 401 sim = F.cosine_similarity(emb1, emb2)[0].item()402 # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).")403 sim_list.append(sim)404 405 return sim_list406