From 8e0edfcf8f2ca8e103165feaae9c4f82ad212dea Mon Sep 17 00:00:00 2001 From: SWivid Date: Thu, 24 Oct 2024 00:55:41 +0800 Subject: [PATCH] final structure. prepared to solve dependencies --- src/f5_tts/{model => eval}/ecapa_tdnn.py | 0 src/f5_tts/eval/utils_eval.py | 397 ++++++++++++++++ src/f5_tts/infer/infer_cli.py | 2 +- src/f5_tts/infer/infer_gradio.py | 6 +- src/f5_tts/infer/speech_edit.py | 2 +- src/f5_tts/{model => infer}/utils_infer.py | 49 +- src/f5_tts/model/utils.py | 438 ------------------ .../datasets}/prepare_csv_wavs.py | 276 +++++------ .../datasets}/prepare_emilia.py | 0 .../datasets}/prepare_wenetspeech4tts.py | 0 10 files changed, 587 insertions(+), 583 deletions(-) rename src/f5_tts/{model => eval}/ecapa_tdnn.py (100%) create mode 100644 src/f5_tts/eval/utils_eval.py rename src/f5_tts/{model => infer}/utils_infer.py (89%) rename src/f5_tts/{scripts => train/datasets}/prepare_csv_wavs.py (97%) rename src/f5_tts/{scripts => train/datasets}/prepare_emilia.py (100%) rename src/f5_tts/{scripts => train/datasets}/prepare_wenetspeech4tts.py (100%) diff --git a/src/f5_tts/model/ecapa_tdnn.py b/src/f5_tts/eval/ecapa_tdnn.py similarity index 100% rename from src/f5_tts/model/ecapa_tdnn.py rename to src/f5_tts/eval/ecapa_tdnn.py diff --git a/src/f5_tts/eval/utils_eval.py b/src/f5_tts/eval/utils_eval.py new file mode 100644 index 0000000..c2cf38e --- /dev/null +++ b/src/f5_tts/eval/utils_eval.py @@ -0,0 +1,397 @@ +import math +import os +import random +import string +from tqdm import tqdm + +import torch +import torch.nn.functional as F +import torchaudio + +from f5_tts.model.modules import MelSpec +from f5_tts.model.utils import convert_char_to_pinyin +from f5_tts.eval.ecapa_tdnn import ECAPA_TDNN_SMALL + + +# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav +def get_seedtts_testset_metainfo(metalst): + f = open(metalst) + lines = f.readlines() + f.close() + metainfo = [] + for line in lines: + if len(line.strip().split("|")) == 5: + utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|") + elif len(line.strip().split("|")) == 4: + utt, prompt_text, prompt_wav, gt_text = line.strip().split("|") + gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav") + if not os.path.isabs(prompt_wav): + prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav) + metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav)) + return metainfo + + +# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav +def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path): + f = open(metalst) + lines = f.readlines() + f.close() + metainfo = [] + for line in lines: + ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t") + + # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc) + ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-") + ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac") + + # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc) + gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-") + gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac") + + metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav)) + + return metainfo + + +# padded to max length mel batch +def padded_mel_batch(ref_mels): + max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax() + padded_ref_mels = [] + for mel in ref_mels: + padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value=0) + padded_ref_mels.append(padded_ref_mel) + padded_ref_mels = torch.stack(padded_ref_mels) + padded_ref_mels = padded_ref_mels.permute(0, 2, 1) + return padded_ref_mels + + +# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav + + +def get_inference_prompt( + metainfo, + speed=1.0, + tokenizer="pinyin", + polyphone=True, + target_sample_rate=24000, + n_mel_channels=100, + hop_length=256, + target_rms=0.1, + use_truth_duration=False, + infer_batch_size=1, + num_buckets=200, + min_secs=3, + max_secs=40, +): + prompts_all = [] + + min_tokens = min_secs * target_sample_rate // hop_length + max_tokens = max_secs * target_sample_rate // hop_length + + batch_accum = [0] * num_buckets + utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = ( + [[] for _ in range(num_buckets)] for _ in range(6) + ) + + mel_spectrogram = MelSpec( + target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length + ) + + for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."): + # Audio + ref_audio, ref_sr = torchaudio.load(prompt_wav) + ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio))) + if ref_rms < target_rms: + ref_audio = ref_audio * target_rms / ref_rms + assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue." + if ref_sr != target_sample_rate: + resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate) + ref_audio = resampler(ref_audio) + + # Text + if len(prompt_text[-1].encode("utf-8")) == 1: + prompt_text = prompt_text + " " + text = [prompt_text + gt_text] + if tokenizer == "pinyin": + text_list = convert_char_to_pinyin(text, polyphone=polyphone) + else: + text_list = text + + # Duration, mel frame length + ref_mel_len = ref_audio.shape[-1] // hop_length + if use_truth_duration: + gt_audio, gt_sr = torchaudio.load(gt_wav) + if gt_sr != target_sample_rate: + resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate) + gt_audio = resampler(gt_audio) + total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed) + + # # test vocoder resynthesis + # ref_audio = gt_audio + else: + ref_text_len = len(prompt_text.encode("utf-8")) + gen_text_len = len(gt_text.encode("utf-8")) + total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed) + + # to mel spectrogram + ref_mel = mel_spectrogram(ref_audio) + ref_mel = ref_mel.squeeze(0) + + # deal with batch + assert infer_batch_size > 0, "infer_batch_size should be greater than 0." + assert ( + min_tokens <= total_mel_len <= max_tokens + ), f"Audio {utt} has duration {total_mel_len*hop_length//target_sample_rate}s out of range [{min_secs}, {max_secs}]." + bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets) + + utts[bucket_i].append(utt) + ref_rms_list[bucket_i].append(ref_rms) + ref_mels[bucket_i].append(ref_mel) + ref_mel_lens[bucket_i].append(ref_mel_len) + total_mel_lens[bucket_i].append(total_mel_len) + final_text_list[bucket_i].extend(text_list) + + batch_accum[bucket_i] += total_mel_len + + if batch_accum[bucket_i] >= infer_batch_size: + # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}") + prompts_all.append( + ( + utts[bucket_i], + ref_rms_list[bucket_i], + padded_mel_batch(ref_mels[bucket_i]), + ref_mel_lens[bucket_i], + total_mel_lens[bucket_i], + final_text_list[bucket_i], + ) + ) + batch_accum[bucket_i] = 0 + ( + utts[bucket_i], + ref_rms_list[bucket_i], + ref_mels[bucket_i], + ref_mel_lens[bucket_i], + total_mel_lens[bucket_i], + final_text_list[bucket_i], + ) = [], [], [], [], [], [] + + # add residual + for bucket_i, bucket_frames in enumerate(batch_accum): + if bucket_frames > 0: + prompts_all.append( + ( + utts[bucket_i], + ref_rms_list[bucket_i], + padded_mel_batch(ref_mels[bucket_i]), + ref_mel_lens[bucket_i], + total_mel_lens[bucket_i], + final_text_list[bucket_i], + ) + ) + # not only leave easy work for last workers + random.seed(666) + random.shuffle(prompts_all) + + return prompts_all + + +# get wav_res_ref_text of seed-tts test metalst +# https://github.com/BytedanceSpeech/seed-tts-eval + + +def get_seed_tts_test(metalst, gen_wav_dir, gpus): + f = open(metalst) + lines = f.readlines() + f.close() + + test_set_ = [] + for line in tqdm(lines): + if len(line.strip().split("|")) == 5: + utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|") + elif len(line.strip().split("|")) == 4: + utt, prompt_text, prompt_wav, gt_text = line.strip().split("|") + + if not os.path.exists(os.path.join(gen_wav_dir, utt + ".wav")): + continue + gen_wav = os.path.join(gen_wav_dir, utt + ".wav") + if not os.path.isabs(prompt_wav): + prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav) + + test_set_.append((gen_wav, prompt_wav, gt_text)) + + num_jobs = len(gpus) + if num_jobs == 1: + return [(gpus[0], test_set_)] + + wav_per_job = len(test_set_) // num_jobs + 1 + test_set = [] + for i in range(num_jobs): + test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job])) + + return test_set + + +# get librispeech test-clean cross sentence test + + +def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth=False): + f = open(metalst) + lines = f.readlines() + f.close() + + test_set_ = [] + for line in tqdm(lines): + ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t") + + if eval_ground_truth: + gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-") + gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac") + else: + if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + ".wav")): + raise FileNotFoundError(f"Generated wav not found: {gen_utt}") + gen_wav = os.path.join(gen_wav_dir, gen_utt + ".wav") + + ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-") + ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac") + + test_set_.append((gen_wav, ref_wav, gen_txt)) + + num_jobs = len(gpus) + if num_jobs == 1: + return [(gpus[0], test_set_)] + + wav_per_job = len(test_set_) // num_jobs + 1 + test_set = [] + for i in range(num_jobs): + test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job])) + + return test_set + + +# load asr model + + +def load_asr_model(lang, ckpt_dir=""): + if lang == "zh": + from funasr import AutoModel + + model = AutoModel( + model=os.path.join(ckpt_dir, "paraformer-zh"), + # vad_model = os.path.join(ckpt_dir, "fsmn-vad"), + # punc_model = os.path.join(ckpt_dir, "ct-punc"), + # spk_model = os.path.join(ckpt_dir, "cam++"), + disable_update=True, + ) # following seed-tts setting + elif lang == "en": + from faster_whisper import WhisperModel + + model_size = "large-v3" if ckpt_dir == "" else ckpt_dir + model = WhisperModel(model_size, device="cuda", compute_type="float16") + return model + + +# WER Evaluation, the way Seed-TTS does + + +def run_asr_wer(args): + rank, lang, test_set, ckpt_dir = args + + if lang == "zh": + import zhconv + + torch.cuda.set_device(rank) + elif lang == "en": + os.environ["CUDA_VISIBLE_DEVICES"] = str(rank) + else: + raise NotImplementedError( + "lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now." + ) + + asr_model = load_asr_model(lang, ckpt_dir=ckpt_dir) + + from zhon.hanzi import punctuation + + punctuation_all = punctuation + string.punctuation + wers = [] + + from jiwer import compute_measures + + for gen_wav, prompt_wav, truth in tqdm(test_set): + if lang == "zh": + res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True) + hypo = res[0]["text"] + hypo = zhconv.convert(hypo, "zh-cn") + elif lang == "en": + segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en") + hypo = "" + for segment in segments: + hypo = hypo + " " + segment.text + + # raw_truth = truth + # raw_hypo = hypo + + for x in punctuation_all: + truth = truth.replace(x, "") + hypo = hypo.replace(x, "") + + truth = truth.replace(" ", " ") + hypo = hypo.replace(" ", " ") + + if lang == "zh": + truth = " ".join([x for x in truth]) + hypo = " ".join([x for x in hypo]) + elif lang == "en": + truth = truth.lower() + hypo = hypo.lower() + + measures = compute_measures(truth, hypo) + wer = measures["wer"] + + # ref_list = truth.split(" ") + # subs = measures["substitutions"] / len(ref_list) + # dele = measures["deletions"] / len(ref_list) + # inse = measures["insertions"] / len(ref_list) + + wers.append(wer) + + return wers + + +# SIM Evaluation + + +def run_sim(args): + rank, test_set, ckpt_dir = args + device = f"cuda:{rank}" + + model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type="wavlm_large", config_path=None) + state_dict = torch.load(ckpt_dir, weights_only=True, map_location=lambda storage, loc: storage) + model.load_state_dict(state_dict["model"], strict=False) + + use_gpu = True if torch.cuda.is_available() else False + if use_gpu: + model = model.cuda(device) + model.eval() + + sim_list = [] + for wav1, wav2, truth in tqdm(test_set): + wav1, sr1 = torchaudio.load(wav1) + wav2, sr2 = torchaudio.load(wav2) + + resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000) + resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000) + wav1 = resample1(wav1) + wav2 = resample2(wav2) + + if use_gpu: + wav1 = wav1.cuda(device) + wav2 = wav2.cuda(device) + with torch.no_grad(): + emb1 = model(wav1) + emb2 = model(wav2) + + sim = F.cosine_similarity(emb1, emb2)[0].item() + # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).") + sim_list.append(sim) + + return sim_list diff --git a/src/f5_tts/infer/infer_cli.py b/src/f5_tts/infer/infer_cli.py index 8cc8ede..546428b 100644 --- a/src/f5_tts/infer/infer_cli.py +++ b/src/f5_tts/infer/infer_cli.py @@ -11,7 +11,7 @@ import tomli from cached_path import cached_path from f5_tts.model import DiT, UNetT -from f5_tts.model.utils_infer import ( +from f5_tts.infer.utils_infer import ( load_vocoder, load_model, preprocess_ref_audio_text, diff --git a/src/f5_tts/infer/infer_gradio.py b/src/f5_tts/infer/infer_gradio.py index 24fbb06..c7fd443 100644 --- a/src/f5_tts/infer/infer_gradio.py +++ b/src/f5_tts/infer/infer_gradio.py @@ -28,15 +28,13 @@ def gpu_decorator(func): from f5_tts.model import DiT, UNetT -from f5_tts.model.utils import ( - save_spectrogram, -) -from f5_tts.model.utils_infer import ( +from f5_tts.infer.utils_infer import ( load_vocoder, load_model, preprocess_ref_audio_text, infer_process, remove_silence_for_generated_wav, + save_spectrogram, ) vocos = load_vocoder() diff --git a/src/f5_tts/infer/speech_edit.py b/src/f5_tts/infer/speech_edit.py index 85b409f..c51f0d1 100644 --- a/src/f5_tts/infer/speech_edit.py +++ b/src/f5_tts/infer/speech_edit.py @@ -10,8 +10,8 @@ from f5_tts.model.utils import ( load_checkpoint, get_tokenizer, convert_char_to_pinyin, - save_spectrogram, ) +from f5_tts.infer.utils_infer import save_spectrogram device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" diff --git a/src/f5_tts/model/utils_infer.py b/src/f5_tts/infer/utils_infer.py similarity index 89% rename from src/f5_tts/model/utils_infer.py rename to src/f5_tts/infer/utils_infer.py index adfdaad..cd8be0d 100644 --- a/src/f5_tts/model/utils_infer.py +++ b/src/f5_tts/infer/utils_infer.py @@ -4,6 +4,11 @@ import re import tempfile +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pylab as plt import numpy as np import torch import torchaudio @@ -14,7 +19,6 @@ from vocos import Vocos from f5_tts.model import CFM from f5_tts.model.utils import ( - load_checkpoint, get_tokenizer, convert_char_to_pinyin, ) @@ -104,6 +108,38 @@ def initialize_asr_pipeline(device=device): ) +# load model checkpoint for inference + + +def load_checkpoint(model, ckpt_path, device, use_ema=True): + if device == "cuda": + model = model.half() + + ckpt_type = ckpt_path.split(".")[-1] + if ckpt_type == "safetensors": + from safetensors.torch import load_file + + checkpoint = load_file(ckpt_path) + else: + checkpoint = torch.load(ckpt_path, weights_only=True) + + if use_ema: + if ckpt_type == "safetensors": + checkpoint = {"ema_model_state_dict": checkpoint} + checkpoint["model_state_dict"] = { + k.replace("ema_model.", ""): v + for k, v in checkpoint["ema_model_state_dict"].items() + if k not in ["initted", "step"] + } + model.load_state_dict(checkpoint["model_state_dict"]) + else: + if ckpt_type == "safetensors": + checkpoint = {"model_state_dict": checkpoint} + model.load_state_dict(checkpoint["model_state_dict"]) + + return model.to(device) + + # load model for inference @@ -355,3 +391,14 @@ def remove_silence_for_generated_wav(filename): non_silent_wave += non_silent_seg aseg = non_silent_wave aseg.export(filename, format="wav") + + +# save spectrogram + + +def save_spectrogram(spectrogram, path): + plt.figure(figsize=(12, 4)) + plt.imshow(spectrogram, origin="lower", aspect="auto") + plt.colorbar() + plt.savefig(path) + plt.close() diff --git a/src/f5_tts/model/utils.py b/src/f5_tts/model/utils.py index bae14e6..a1e544c 100644 --- a/src/f5_tts/model/utils.py +++ b/src/f5_tts/model/utils.py @@ -1,29 +1,16 @@ from __future__ import annotations import os -import math import random -import string from importlib.resources import files -from tqdm import tqdm from collections import defaultdict -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pylab as plt - import torch -import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence -import torchaudio import jieba from pypinyin import lazy_pinyin, Style -from f5_tts.model.ecapa_tdnn import ECAPA_TDNN_SMALL -from f5_tts.model.modules import MelSpec - # seed everything @@ -183,399 +170,6 @@ def convert_char_to_pinyin(text_list, polyphone=True): return final_text_list -# save spectrogram -def save_spectrogram(spectrogram, path): - plt.figure(figsize=(12, 4)) - plt.imshow(spectrogram, origin="lower", aspect="auto") - plt.colorbar() - plt.savefig(path) - plt.close() - - -# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav -def get_seedtts_testset_metainfo(metalst): - f = open(metalst) - lines = f.readlines() - f.close() - metainfo = [] - for line in lines: - if len(line.strip().split("|")) == 5: - utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|") - elif len(line.strip().split("|")) == 4: - utt, prompt_text, prompt_wav, gt_text = line.strip().split("|") - gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav") - if not os.path.isabs(prompt_wav): - prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav) - metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav)) - return metainfo - - -# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav -def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path): - f = open(metalst) - lines = f.readlines() - f.close() - metainfo = [] - for line in lines: - ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t") - - # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc) - ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-") - ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac") - - # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc) - gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-") - gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac") - - metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav)) - - return metainfo - - -# padded to max length mel batch -def padded_mel_batch(ref_mels): - max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax() - padded_ref_mels = [] - for mel in ref_mels: - padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value=0) - padded_ref_mels.append(padded_ref_mel) - padded_ref_mels = torch.stack(padded_ref_mels) - padded_ref_mels = padded_ref_mels.permute(0, 2, 1) - return padded_ref_mels - - -# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav - - -def get_inference_prompt( - metainfo, - speed=1.0, - tokenizer="pinyin", - polyphone=True, - target_sample_rate=24000, - n_mel_channels=100, - hop_length=256, - target_rms=0.1, - use_truth_duration=False, - infer_batch_size=1, - num_buckets=200, - min_secs=3, - max_secs=40, -): - prompts_all = [] - - min_tokens = min_secs * target_sample_rate // hop_length - max_tokens = max_secs * target_sample_rate // hop_length - - batch_accum = [0] * num_buckets - utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = ( - [[] for _ in range(num_buckets)] for _ in range(6) - ) - - mel_spectrogram = MelSpec( - target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length - ) - - for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."): - # Audio - ref_audio, ref_sr = torchaudio.load(prompt_wav) - ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio))) - if ref_rms < target_rms: - ref_audio = ref_audio * target_rms / ref_rms - assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue." - if ref_sr != target_sample_rate: - resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate) - ref_audio = resampler(ref_audio) - - # Text - if len(prompt_text[-1].encode("utf-8")) == 1: - prompt_text = prompt_text + " " - text = [prompt_text + gt_text] - if tokenizer == "pinyin": - text_list = convert_char_to_pinyin(text, polyphone=polyphone) - else: - text_list = text - - # Duration, mel frame length - ref_mel_len = ref_audio.shape[-1] // hop_length - if use_truth_duration: - gt_audio, gt_sr = torchaudio.load(gt_wav) - if gt_sr != target_sample_rate: - resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate) - gt_audio = resampler(gt_audio) - total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed) - - # # test vocoder resynthesis - # ref_audio = gt_audio - else: - ref_text_len = len(prompt_text.encode("utf-8")) - gen_text_len = len(gt_text.encode("utf-8")) - total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed) - - # to mel spectrogram - ref_mel = mel_spectrogram(ref_audio) - ref_mel = ref_mel.squeeze(0) - - # deal with batch - assert infer_batch_size > 0, "infer_batch_size should be greater than 0." - assert ( - min_tokens <= total_mel_len <= max_tokens - ), f"Audio {utt} has duration {total_mel_len*hop_length//target_sample_rate}s out of range [{min_secs}, {max_secs}]." - bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets) - - utts[bucket_i].append(utt) - ref_rms_list[bucket_i].append(ref_rms) - ref_mels[bucket_i].append(ref_mel) - ref_mel_lens[bucket_i].append(ref_mel_len) - total_mel_lens[bucket_i].append(total_mel_len) - final_text_list[bucket_i].extend(text_list) - - batch_accum[bucket_i] += total_mel_len - - if batch_accum[bucket_i] >= infer_batch_size: - # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}") - prompts_all.append( - ( - utts[bucket_i], - ref_rms_list[bucket_i], - padded_mel_batch(ref_mels[bucket_i]), - ref_mel_lens[bucket_i], - total_mel_lens[bucket_i], - final_text_list[bucket_i], - ) - ) - batch_accum[bucket_i] = 0 - ( - utts[bucket_i], - ref_rms_list[bucket_i], - ref_mels[bucket_i], - ref_mel_lens[bucket_i], - total_mel_lens[bucket_i], - final_text_list[bucket_i], - ) = [], [], [], [], [], [] - - # add residual - for bucket_i, bucket_frames in enumerate(batch_accum): - if bucket_frames > 0: - prompts_all.append( - ( - utts[bucket_i], - ref_rms_list[bucket_i], - padded_mel_batch(ref_mels[bucket_i]), - ref_mel_lens[bucket_i], - total_mel_lens[bucket_i], - final_text_list[bucket_i], - ) - ) - # not only leave easy work for last workers - random.seed(666) - random.shuffle(prompts_all) - - return prompts_all - - -# get wav_res_ref_text of seed-tts test metalst -# https://github.com/BytedanceSpeech/seed-tts-eval - - -def get_seed_tts_test(metalst, gen_wav_dir, gpus): - f = open(metalst) - lines = f.readlines() - f.close() - - test_set_ = [] - for line in tqdm(lines): - if len(line.strip().split("|")) == 5: - utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|") - elif len(line.strip().split("|")) == 4: - utt, prompt_text, prompt_wav, gt_text = line.strip().split("|") - - if not os.path.exists(os.path.join(gen_wav_dir, utt + ".wav")): - continue - gen_wav = os.path.join(gen_wav_dir, utt + ".wav") - if not os.path.isabs(prompt_wav): - prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav) - - test_set_.append((gen_wav, prompt_wav, gt_text)) - - num_jobs = len(gpus) - if num_jobs == 1: - return [(gpus[0], test_set_)] - - wav_per_job = len(test_set_) // num_jobs + 1 - test_set = [] - for i in range(num_jobs): - test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job])) - - return test_set - - -# get librispeech test-clean cross sentence test - - -def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth=False): - f = open(metalst) - lines = f.readlines() - f.close() - - test_set_ = [] - for line in tqdm(lines): - ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t") - - if eval_ground_truth: - gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-") - gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac") - else: - if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + ".wav")): - raise FileNotFoundError(f"Generated wav not found: {gen_utt}") - gen_wav = os.path.join(gen_wav_dir, gen_utt + ".wav") - - ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-") - ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac") - - test_set_.append((gen_wav, ref_wav, gen_txt)) - - num_jobs = len(gpus) - if num_jobs == 1: - return [(gpus[0], test_set_)] - - wav_per_job = len(test_set_) // num_jobs + 1 - test_set = [] - for i in range(num_jobs): - test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job])) - - return test_set - - -# load asr model - - -def load_asr_model(lang, ckpt_dir=""): - if lang == "zh": - from funasr import AutoModel - - model = AutoModel( - model=os.path.join(ckpt_dir, "paraformer-zh"), - # vad_model = os.path.join(ckpt_dir, "fsmn-vad"), - # punc_model = os.path.join(ckpt_dir, "ct-punc"), - # spk_model = os.path.join(ckpt_dir, "cam++"), - disable_update=True, - ) # following seed-tts setting - elif lang == "en": - from faster_whisper import WhisperModel - - model_size = "large-v3" if ckpt_dir == "" else ckpt_dir - model = WhisperModel(model_size, device="cuda", compute_type="float16") - return model - - -# WER Evaluation, the way Seed-TTS does - - -def run_asr_wer(args): - rank, lang, test_set, ckpt_dir = args - - if lang == "zh": - import zhconv - - torch.cuda.set_device(rank) - elif lang == "en": - os.environ["CUDA_VISIBLE_DEVICES"] = str(rank) - else: - raise NotImplementedError( - "lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now." - ) - - asr_model = load_asr_model(lang, ckpt_dir=ckpt_dir) - - from zhon.hanzi import punctuation - - punctuation_all = punctuation + string.punctuation - wers = [] - - from jiwer import compute_measures - - for gen_wav, prompt_wav, truth in tqdm(test_set): - if lang == "zh": - res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True) - hypo = res[0]["text"] - hypo = zhconv.convert(hypo, "zh-cn") - elif lang == "en": - segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en") - hypo = "" - for segment in segments: - hypo = hypo + " " + segment.text - - # raw_truth = truth - # raw_hypo = hypo - - for x in punctuation_all: - truth = truth.replace(x, "") - hypo = hypo.replace(x, "") - - truth = truth.replace(" ", " ") - hypo = hypo.replace(" ", " ") - - if lang == "zh": - truth = " ".join([x for x in truth]) - hypo = " ".join([x for x in hypo]) - elif lang == "en": - truth = truth.lower() - hypo = hypo.lower() - - measures = compute_measures(truth, hypo) - wer = measures["wer"] - - # ref_list = truth.split(" ") - # subs = measures["substitutions"] / len(ref_list) - # dele = measures["deletions"] / len(ref_list) - # inse = measures["insertions"] / len(ref_list) - - wers.append(wer) - - return wers - - -# SIM Evaluation - - -def run_sim(args): - rank, test_set, ckpt_dir = args - device = f"cuda:{rank}" - - model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type="wavlm_large", config_path=None) - state_dict = torch.load(ckpt_dir, weights_only=True, map_location=lambda storage, loc: storage) - model.load_state_dict(state_dict["model"], strict=False) - - use_gpu = True if torch.cuda.is_available() else False - if use_gpu: - model = model.cuda(device) - model.eval() - - sim_list = [] - for wav1, wav2, truth in tqdm(test_set): - wav1, sr1 = torchaudio.load(wav1) - wav2, sr2 = torchaudio.load(wav2) - - resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000) - resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000) - wav1 = resample1(wav1) - wav2 = resample2(wav2) - - if use_gpu: - wav1 = wav1.cuda(device) - wav2 = wav2.cuda(device) - with torch.no_grad(): - emb1 = model(wav1) - emb2 = model(wav2) - - sim = F.cosine_similarity(emb1, emb2)[0].item() - # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).") - sim_list.append(sim) - - return sim_list - - # filter func for dirty data with many repetitions @@ -588,35 +182,3 @@ def repetition_found(text, length=2, tolerance=10): if count > tolerance: return True return False - - -# load model checkpoint for inference - - -def load_checkpoint(model, ckpt_path, device, use_ema=True): - if device == "cuda": - model = model.half() - - ckpt_type = ckpt_path.split(".")[-1] - if ckpt_type == "safetensors": - from safetensors.torch import load_file - - checkpoint = load_file(ckpt_path) - else: - checkpoint = torch.load(ckpt_path, weights_only=True) - - if use_ema: - if ckpt_type == "safetensors": - checkpoint = {"ema_model_state_dict": checkpoint} - checkpoint["model_state_dict"] = { - k.replace("ema_model.", ""): v - for k, v in checkpoint["ema_model_state_dict"].items() - if k not in ["initted", "step"] - } - model.load_state_dict(checkpoint["model_state_dict"]) - else: - if ckpt_type == "safetensors": - checkpoint = {"model_state_dict": checkpoint} - model.load_state_dict(checkpoint["model_state_dict"]) - - return model.to(device) diff --git a/src/f5_tts/scripts/prepare_csv_wavs.py b/src/f5_tts/train/datasets/prepare_csv_wavs.py similarity index 97% rename from src/f5_tts/scripts/prepare_csv_wavs.py rename to src/f5_tts/train/datasets/prepare_csv_wavs.py index 7fba925..e68b053 100644 --- a/src/f5_tts/scripts/prepare_csv_wavs.py +++ b/src/f5_tts/train/datasets/prepare_csv_wavs.py @@ -1,138 +1,138 @@ -import sys -import os - -sys.path.append(os.getcwd()) - -from pathlib import Path -import json -import shutil -import argparse - -import csv -import torchaudio -from tqdm import tqdm -from datasets.arrow_writer import ArrowWriter - -from f5_tts.model.utils import ( - convert_char_to_pinyin, -) - -PRETRAINED_VOCAB_PATH = Path(__file__).parent.parent / "data/Emilia_ZH_EN_pinyin/vocab.txt" - - -def is_csv_wavs_format(input_dataset_dir): - fpath = Path(input_dataset_dir) - metadata = fpath / "metadata.csv" - wavs = fpath / "wavs" - return metadata.exists() and metadata.is_file() and wavs.exists() and wavs.is_dir() - - -def prepare_csv_wavs_dir(input_dir): - assert is_csv_wavs_format(input_dir), f"not csv_wavs format: {input_dir}" - input_dir = Path(input_dir) - metadata_path = input_dir / "metadata.csv" - audio_path_text_pairs = read_audio_text_pairs(metadata_path.as_posix()) - - sub_result, durations = [], [] - vocab_set = set() - polyphone = True - for audio_path, text in audio_path_text_pairs: - if not Path(audio_path).exists(): - print(f"audio {audio_path} not found, skipping") - continue - audio_duration = get_audio_duration(audio_path) - # assume tokenizer = "pinyin" ("pinyin" | "char") - text = convert_char_to_pinyin([text], polyphone=polyphone)[0] - sub_result.append({"audio_path": audio_path, "text": text, "duration": audio_duration}) - durations.append(audio_duration) - vocab_set.update(list(text)) - - return sub_result, durations, vocab_set - - -def get_audio_duration(audio_path): - audio, sample_rate = torchaudio.load(audio_path) - num_channels = audio.shape[0] - return audio.shape[1] / (sample_rate * num_channels) - - -def read_audio_text_pairs(csv_file_path): - audio_text_pairs = [] - - parent = Path(csv_file_path).parent - with open(csv_file_path, mode="r", newline="", encoding="utf-8") as csvfile: - reader = csv.reader(csvfile, delimiter="|") - next(reader) # Skip the header row - for row in reader: - if len(row) >= 2: - audio_file = row[0].strip() # First column: audio file path - text = row[1].strip() # Second column: text - audio_file_path = parent / audio_file - audio_text_pairs.append((audio_file_path.as_posix(), text)) - - return audio_text_pairs - - -def save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_finetune): - out_dir = Path(out_dir) - # save preprocessed dataset to disk - out_dir.mkdir(exist_ok=True, parents=True) - print(f"\nSaving to {out_dir} ...") - - # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list}) # oom - # dataset.save_to_disk(f"data/{dataset_name}/raw", max_shard_size="2GB") - raw_arrow_path = out_dir / "raw.arrow" - with ArrowWriter(path=raw_arrow_path.as_posix(), writer_batch_size=1) as writer: - for line in tqdm(result, desc="Writing to raw.arrow ..."): - writer.write(line) - - # dup a json separately saving duration in case for DynamicBatchSampler ease - dur_json_path = out_dir / "duration.json" - with open(dur_json_path.as_posix(), "w", encoding="utf-8") as f: - json.dump({"duration": duration_list}, f, ensure_ascii=False) - - # vocab map, i.e. tokenizer - # add alphabets and symbols (optional, if plan to ft on de/fr etc.) - # if tokenizer == "pinyin": - # text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)]) - voca_out_path = out_dir / "vocab.txt" - with open(voca_out_path.as_posix(), "w") as f: - for vocab in sorted(text_vocab_set): - f.write(vocab + "\n") - - if is_finetune: - file_vocab_finetune = PRETRAINED_VOCAB_PATH.as_posix() - shutil.copy2(file_vocab_finetune, voca_out_path) - else: - with open(voca_out_path, "w") as f: - for vocab in sorted(text_vocab_set): - f.write(vocab + "\n") - - dataset_name = out_dir.stem - print(f"\nFor {dataset_name}, sample count: {len(result)}") - print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") - print(f"For {dataset_name}, total {sum(duration_list)/3600:.2f} hours") - - -def prepare_and_save_set(inp_dir, out_dir, is_finetune: bool = True): - if is_finetune: - assert PRETRAINED_VOCAB_PATH.exists(), f"pretrained vocab.txt not found: {PRETRAINED_VOCAB_PATH}" - sub_result, durations, vocab_set = prepare_csv_wavs_dir(inp_dir) - save_prepped_dataset(out_dir, sub_result, durations, vocab_set, is_finetune) - - -def cli(): - # finetune: python scripts/prepare_csv_wavs.py /path/to/input_dir /path/to/output_dir_pinyin - # pretrain: python scripts/prepare_csv_wavs.py /path/to/output_dir_pinyin --pretrain - parser = argparse.ArgumentParser(description="Prepare and save dataset.") - parser.add_argument("inp_dir", type=str, help="Input directory containing the data.") - parser.add_argument("out_dir", type=str, help="Output directory to save the prepared data.") - parser.add_argument("--pretrain", action="store_true", help="Enable for new pretrain, otherwise is a fine-tune") - - args = parser.parse_args() - - prepare_and_save_set(args.inp_dir, args.out_dir, is_finetune=not args.pretrain) - - -if __name__ == "__main__": - cli() +import sys +import os + +sys.path.append(os.getcwd()) + +from pathlib import Path +import json +import shutil +import argparse + +import csv +import torchaudio +from tqdm import tqdm +from datasets.arrow_writer import ArrowWriter + +from f5_tts.model.utils import ( + convert_char_to_pinyin, +) + +PRETRAINED_VOCAB_PATH = Path(__file__).parent.parent / "data/Emilia_ZH_EN_pinyin/vocab.txt" + + +def is_csv_wavs_format(input_dataset_dir): + fpath = Path(input_dataset_dir) + metadata = fpath / "metadata.csv" + wavs = fpath / "wavs" + return metadata.exists() and metadata.is_file() and wavs.exists() and wavs.is_dir() + + +def prepare_csv_wavs_dir(input_dir): + assert is_csv_wavs_format(input_dir), f"not csv_wavs format: {input_dir}" + input_dir = Path(input_dir) + metadata_path = input_dir / "metadata.csv" + audio_path_text_pairs = read_audio_text_pairs(metadata_path.as_posix()) + + sub_result, durations = [], [] + vocab_set = set() + polyphone = True + for audio_path, text in audio_path_text_pairs: + if not Path(audio_path).exists(): + print(f"audio {audio_path} not found, skipping") + continue + audio_duration = get_audio_duration(audio_path) + # assume tokenizer = "pinyin" ("pinyin" | "char") + text = convert_char_to_pinyin([text], polyphone=polyphone)[0] + sub_result.append({"audio_path": audio_path, "text": text, "duration": audio_duration}) + durations.append(audio_duration) + vocab_set.update(list(text)) + + return sub_result, durations, vocab_set + + +def get_audio_duration(audio_path): + audio, sample_rate = torchaudio.load(audio_path) + num_channels = audio.shape[0] + return audio.shape[1] / (sample_rate * num_channels) + + +def read_audio_text_pairs(csv_file_path): + audio_text_pairs = [] + + parent = Path(csv_file_path).parent + with open(csv_file_path, mode="r", newline="", encoding="utf-8") as csvfile: + reader = csv.reader(csvfile, delimiter="|") + next(reader) # Skip the header row + for row in reader: + if len(row) >= 2: + audio_file = row[0].strip() # First column: audio file path + text = row[1].strip() # Second column: text + audio_file_path = parent / audio_file + audio_text_pairs.append((audio_file_path.as_posix(), text)) + + return audio_text_pairs + + +def save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_finetune): + out_dir = Path(out_dir) + # save preprocessed dataset to disk + out_dir.mkdir(exist_ok=True, parents=True) + print(f"\nSaving to {out_dir} ...") + + # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list}) # oom + # dataset.save_to_disk(f"data/{dataset_name}/raw", max_shard_size="2GB") + raw_arrow_path = out_dir / "raw.arrow" + with ArrowWriter(path=raw_arrow_path.as_posix(), writer_batch_size=1) as writer: + for line in tqdm(result, desc="Writing to raw.arrow ..."): + writer.write(line) + + # dup a json separately saving duration in case for DynamicBatchSampler ease + dur_json_path = out_dir / "duration.json" + with open(dur_json_path.as_posix(), "w", encoding="utf-8") as f: + json.dump({"duration": duration_list}, f, ensure_ascii=False) + + # vocab map, i.e. tokenizer + # add alphabets and symbols (optional, if plan to ft on de/fr etc.) + # if tokenizer == "pinyin": + # text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)]) + voca_out_path = out_dir / "vocab.txt" + with open(voca_out_path.as_posix(), "w") as f: + for vocab in sorted(text_vocab_set): + f.write(vocab + "\n") + + if is_finetune: + file_vocab_finetune = PRETRAINED_VOCAB_PATH.as_posix() + shutil.copy2(file_vocab_finetune, voca_out_path) + else: + with open(voca_out_path, "w") as f: + for vocab in sorted(text_vocab_set): + f.write(vocab + "\n") + + dataset_name = out_dir.stem + print(f"\nFor {dataset_name}, sample count: {len(result)}") + print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") + print(f"For {dataset_name}, total {sum(duration_list)/3600:.2f} hours") + + +def prepare_and_save_set(inp_dir, out_dir, is_finetune: bool = True): + if is_finetune: + assert PRETRAINED_VOCAB_PATH.exists(), f"pretrained vocab.txt not found: {PRETRAINED_VOCAB_PATH}" + sub_result, durations, vocab_set = prepare_csv_wavs_dir(inp_dir) + save_prepped_dataset(out_dir, sub_result, durations, vocab_set, is_finetune) + + +def cli(): + # finetune: python scripts/prepare_csv_wavs.py /path/to/input_dir /path/to/output_dir_pinyin + # pretrain: python scripts/prepare_csv_wavs.py /path/to/output_dir_pinyin --pretrain + parser = argparse.ArgumentParser(description="Prepare and save dataset.") + parser.add_argument("inp_dir", type=str, help="Input directory containing the data.") + parser.add_argument("out_dir", type=str, help="Output directory to save the prepared data.") + parser.add_argument("--pretrain", action="store_true", help="Enable for new pretrain, otherwise is a fine-tune") + + args = parser.parse_args() + + prepare_and_save_set(args.inp_dir, args.out_dir, is_finetune=not args.pretrain) + + +if __name__ == "__main__": + cli() diff --git a/src/f5_tts/scripts/prepare_emilia.py b/src/f5_tts/train/datasets/prepare_emilia.py similarity index 100% rename from src/f5_tts/scripts/prepare_emilia.py rename to src/f5_tts/train/datasets/prepare_emilia.py diff --git a/src/f5_tts/scripts/prepare_wenetspeech4tts.py b/src/f5_tts/train/datasets/prepare_wenetspeech4tts.py similarity index 100% rename from src/f5_tts/scripts/prepare_wenetspeech4tts.py rename to src/f5_tts/train/datasets/prepare_wenetspeech4tts.py