Upload 6 files

model/__init__.py

@@ -1,7 +1,10 @@
-from model.cfm import CFM
+from f5_tts.model.cfm import CFM
 
-from model.backbones.unett import UNetT
-from model.backbones.dit import DiT
-from model.backbones.mmdit import MMDiT
+from f5_tts.model.backbones.unett import UNetT
+from f5_tts.model.backbones.dit import DiT
+from f5_tts.model.backbones.mmdit import MMDiT
 
-from model.trainer import Trainer
+from f5_tts.model.trainer import Trainer
+
+
+__all__ = ["CFM", "UNetT", "DiT", "MMDiT", "Trainer"]

model/cfm.py

@@ -18,8 +18,8 @@ from torch import nn
 from torch.nn.utils.rnn import pad_sequence
 from torchdiffeq import odeint
 
-from model.modules import MelSpec
-from model.utils import (
+from f5_tts.model.modules import MelSpec
+from f5_tts.model.utils import (
     default,
     exists,
     lens_to_mask,
@@ -193,7 +193,7 @@ class CFM(nn.Module):
             y0 = (1 - t_start) * y0 + t_start * test_cond
             steps = int(steps * (1 - t_start))
 
-        t = torch.linspace(t_start, 1, steps, device=self.device, dtype=step_cond.dtype)
+        t = torch.linspace(t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype)
         if sway_sampling_coef is not None:
             t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
 

model/dataset.py

@@ -11,8 +11,8 @@ from torch import nn
 from torch.utils.data import Dataset, Sampler
 from tqdm import tqdm
 
-from model.modules import MelSpec
-from model.utils import default
+from f5_tts.model.modules import MelSpec
+from f5_tts.model.utils import default
 
 
 class HFDataset(Dataset):
@@ -127,38 +127,43 @@ class CustomDataset(Dataset):
         return len(self.data)
 
     def __getitem__(self, index):
-        row = self.data[index]
-        audio_path = row["audio_path"]
-        text = row["text"]
-        duration = row["duration"]
+        while True:
+            row = self.data[index]
+            audio_path = row["audio_path"]
+            text = row["text"]
+            duration = row["duration"]
+
+            # filter by given length
+            if 0.3 <= duration <= 30:
+                break  # valid
+
+            index = (index + 1) % len(self.data)
 
         if self.preprocessed_mel:
             mel_spec = torch.tensor(row["mel_spec"])
-
         else:
             audio, source_sample_rate = torchaudio.load(audio_path)
+
+            # make sure mono input
             if audio.shape[0] > 1:
                 audio = torch.mean(audio, dim=0, keepdim=True)
 
-            if duration > 30 or duration < 0.3:
-                return self.__getitem__((index + 1) % len(self.data))
-
+            # resample if necessary
             if source_sample_rate != self.target_sample_rate:
                 resampler = torchaudio.transforms.Resample(source_sample_rate, self.target_sample_rate)
                 audio = resampler(audio)
 
+            # to mel spectrogram
             mel_spec = self.mel_spectrogram(audio)
-            mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t')
+            mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t'
 
-        return dict(
-            mel_spec=mel_spec,
-            text=text,
-        )
+        return {
+            "mel_spec": mel_spec,
+            "text": text,
+        }
 
 
 # Dynamic Batch Sampler
-
-
 class DynamicBatchSampler(Sampler[list[int]]):
     """Extension of Sampler that will do the following:
     1.  Change the batch size (essentially number of sequences)

model/trainer.py

@@ -14,9 +14,9 @@ from torch.optim.lr_scheduler import LinearLR, SequentialLR
 from torch.utils.data import DataLoader, Dataset, SequentialSampler
 from tqdm import tqdm
 
-from model import CFM
-from model.dataset import DynamicBatchSampler, collate_fn
-from model.utils import default, exists
+from f5_tts.model import CFM
+from f5_tts.model.dataset import DynamicBatchSampler, collate_fn
+from f5_tts.model.utils import default, exists
 
 # trainer
 
@@ -47,6 +47,8 @@ class Trainer:
         ema_kwargs: dict = dict(),
         bnb_optimizer: bool = False,
         mel_spec_type: str = "vocos",  # "vocos" | "bigvgan"
+        is_local_vocoder: bool = False,  # use local path vocoder
+        local_vocoder_path: str = "",  # local vocoder path
     ):
         ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
 
@@ -108,7 +110,11 @@ class Trainer:
         self.max_samples = max_samples
         self.grad_accumulation_steps = grad_accumulation_steps
         self.max_grad_norm = max_grad_norm
+
+        # mel vocoder config
         self.vocoder_name = mel_spec_type
+        self.is_local_vocoder = is_local_vocoder
+        self.local_vocoder_path = local_vocoder_path
 
         self.noise_scheduler = noise_scheduler
 
@@ -148,7 +154,7 @@ class Trainer:
         if (
             not exists(self.checkpoint_path)
             or not os.path.exists(self.checkpoint_path)
-            or not os.listdir(self.checkpoint_path)
+            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
         ):
             return 0
 
@@ -199,7 +205,9 @@ class Trainer:
         if self.log_samples:
             from f5_tts.infer.utils_infer import cfg_strength, load_vocoder, nfe_step, sway_sampling_coef
 
-            vocoder = load_vocoder(vocoder_name=self.vocoder_name)
+            vocoder = load_vocoder(
+                vocoder_name=self.vocoder_name, is_local=self.is_local_vocoder, local_path=self.local_vocoder_path
+            )
             target_sample_rate = self.accelerator.unwrap_model(self.model).mel_spec.target_sample_rate
             log_samples_path = f"{self.checkpoint_path}/samples"
             os.makedirs(log_samples_path, exist_ok=True)
@@ -324,26 +332,31 @@ class Trainer:
                     self.save_checkpoint(global_step)
 
                     if self.log_samples and self.accelerator.is_local_main_process:
-                        ref_audio, ref_audio_len = vocoder.decode(batch["mel"][0].unsqueeze(0)), mel_lengths[0]
-                        torchaudio.save(
-                            f"{log_samples_path}/step_{global_step}_ref.wav", ref_audio.cpu(), target_sample_rate
-                        )
+                        ref_audio_len = mel_lengths[0]
+                        infer_text = [
+                            text_inputs[0] + ([" "] if isinstance(text_inputs[0], list) else " ") + text_inputs[0]
+                        ]
                         with torch.inference_mode():
                             generated, _ = self.accelerator.unwrap_model(self.model).sample(
                                 cond=mel_spec[0][:ref_audio_len].unsqueeze(0),
-                                text=[text_inputs[0] + [" "] + text_inputs[0]],
+                                text=infer_text,
                                 duration=ref_audio_len * 2,
                                 steps=nfe_step,
                                 cfg_strength=cfg_strength,
                                 sway_sampling_coef=sway_sampling_coef,
                             )
-                        generated = generated.to(torch.float32)
-                        gen_audio = vocoder.decode(
-                            generated[:, ref_audio_len:, :].permute(0, 2, 1).to(self.accelerator.device)
-                        )
-                        torchaudio.save(
-                            f"{log_samples_path}/step_{global_step}_gen.wav", gen_audio.cpu(), target_sample_rate
-                        )
+                            generated = generated.to(torch.float32)
+                            gen_mel_spec = generated[:, ref_audio_len:, :].permute(0, 2, 1).to(self.accelerator.device)
+                            ref_mel_spec = batch["mel"][0].unsqueeze(0)
+                            if self.vocoder_name == "vocos":
+                                gen_audio = vocoder.decode(gen_mel_spec).cpu()
+                                ref_audio = vocoder.decode(ref_mel_spec).cpu()
+                            elif self.vocoder_name == "bigvgan":
+                                gen_audio = vocoder(gen_mel_spec).squeeze(0).cpu()
+                                ref_audio = vocoder(ref_mel_spec).squeeze(0).cpu()
+
+                        torchaudio.save(f"{log_samples_path}/step_{global_step}_gen.wav", gen_audio, target_sample_rate)
+                        torchaudio.save(f"{log_samples_path}/step_{global_step}_ref.wav", ref_audio, target_sample_rate)
 
                 if global_step % self.last_per_steps == 0:
                     self.save_checkpoint(global_step, last=True)

model/utils.py

@@ -1,135 +1,116 @@
 from __future__ import annotations
 
 import os
-import re
-import math
 import random
-import string
-from tqdm import tqdm
 from collections import defaultdict
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pylab as plt
+from importlib.resources import files
 
 import torch
-import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
-import torchaudio
-
-import einx
-from einops import rearrange, reduce
 
 import jieba
 from pypinyin import lazy_pinyin, Style
 
-from model.ecapa_tdnn import ECAPA_TDNN_SMALL
-from model.modules import MelSpec
-
 
 # seed everything
 
-def seed_everything(seed = 0):
+
+def seed_everything(seed=0):
     random.seed(seed)
-    os.environ['PYTHONHASHSEED'] = str(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
+
 # helpers
 
+
 def exists(v):
     return v is not None
 
+
 def default(v, d):
     return v if exists(v) else d
 
+
 # tensor helpers
 
-def lens_to_mask(
-    t: int['b'],
-    length: int | None = None
-) -> bool['b n']:
 
+def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
     if not exists(length):
         length = t.amax()
 
-    seq = torch.arange(length, device = t.device)
-    return einx.less('n, b -> b n', seq, t)
-
-def mask_from_start_end_indices(
-    seq_len: int['b'],
-    start: int['b'],
-    end: int['b']
-):
-    max_seq_len = seq_len.max().item()  
-    seq = torch.arange(max_seq_len, device = start.device).long()
-    return einx.greater_equal('n, b -> b n', seq, start) & einx.less('n, b -> b n', seq, end)
-
-def mask_from_frac_lengths(
-    seq_len: int['b'],
-    frac_lengths: float['b']
-):
+    seq = torch.arange(length, device=t.device)
+    return seq[None, :] < t[:, None]
+
+
+def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F821
+    max_seq_len = seq_len.max().item()
+    seq = torch.arange(max_seq_len, device=start.device).long()
+    start_mask = seq[None, :] >= start[:, None]
+    end_mask = seq[None, :] < end[:, None]
+    return start_mask & end_mask
+
+
+def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F821
     lengths = (frac_lengths * seq_len).long()
     max_start = seq_len - lengths
 
     rand = torch.rand_like(frac_lengths)
-    start = (max_start * rand).long().clamp(min = 0)
+    start = (max_start * rand).long().clamp(min=0)
     end = start + lengths
 
     return mask_from_start_end_indices(seq_len, start, end)
 
-def maybe_masked_mean(
-    t: float['b n d'],
-    mask: bool['b n'] = None
-) -> float['b d']:
 
+def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F722
     if not exists(mask):
-        return t.mean(dim = 1)
+        return t.mean(dim=1)
 
-    t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
-    num = reduce(t, 'b n d -> b d', 'sum')
-    den = reduce(mask.float(), 'b n -> b', 'sum')
+    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))
+    num = t.sum(dim=1)
+    den = mask.float().sum(dim=1)
 
-    return einx.divide('b d, b -> b d', num, den.clamp(min = 1.))
+    return num / den.clamp(min=1.0)
 
 
 # simple utf-8 tokenizer, since paper went character based
-def list_str_to_tensor(
-    text: list[str],
-    padding_value = -1
-) -> int['b nt']:
-    list_tensors = [torch.tensor([*bytes(t, 'UTF-8')]) for t in text]  # ByT5 style
-    text = pad_sequence(list_tensors, padding_value = padding_value, batch_first = True)
+def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
+    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style
+    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)
     return text
 
+
 # char tokenizer, based on custom dataset's extracted .txt file
 def list_str_to_idx(
     text: list[str] | list[list[str]],
     vocab_char_map: dict[str, int],  # {char: idx}
-    padding_value = -1
-) -> int['b nt']:
+    padding_value=-1,
+) -> int["b nt"]:  # noqa: F722
     list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
-    text = pad_sequence(list_idx_tensors, padding_value = padding_value, batch_first = True)
+    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
     return text
 
 
 # Get tokenizer
 
+
 def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
-    ''' 
+    """
     tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
                 - "char" for char-wise tokenizer, need .txt vocab_file
                 - "byte" for utf-8 tokenizer
                 - "custom" if you're directly passing in a path to the vocab.txt you want to use
     vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
                 - if use "char", derived from unfiltered character & symbol counts of custom dataset
-                - if use "byte", set to 256 (unicode byte range) 
-    ''' 
+                - if use "byte", set to 256 (unicode byte range)
+    """
     if tokenizer in ["pinyin", "char"]:
-        with open (f"data/{dataset_name}_{tokenizer}/vocab.txt", "r", encoding="utf-8") as f:
+        tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
+        with open(tokenizer_path, "r", encoding="utf-8") as f:
             vocab_char_map = {}
             for i, char in enumerate(f):
                 vocab_char_map[char[:-1]] = i
@@ -139,8 +120,9 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
     elif tokenizer == "byte":
         vocab_char_map = None
         vocab_size = 256
+
     elif tokenizer == "custom":
-        with open (dataset_name, "r", encoding="utf-8") as f:
+        with open(dataset_name, "r", encoding="utf-8") as f:
             vocab_char_map = {}
             for i, char in enumerate(f):
                 vocab_char_map[char[:-1]] = i
@@ -151,16 +133,19 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
 
 # convert char to pinyin
 
-def convert_char_to_pinyin(text_list, polyphone = True):
+
+def convert_char_to_pinyin(text_list, polyphone=True):
     final_text_list = []
-    god_knows_why_en_testset_contains_zh_quote = str.maketrans({'“': '"', '”': '"', '‘': "'", '’': "'"})  # in case librispeech (orig no-pc) test-clean
-    custom_trans = str.maketrans({';': ','})  # add custom trans here, to address oov
+    god_knows_why_en_testset_contains_zh_quote = str.maketrans(
+        {"“": '"', "”": '"', "‘": "'", "’": "'"}
+    )  # in case librispeech (orig no-pc) test-clean
+    custom_trans = str.maketrans({";": ","})  # add custom trans here, to address oov
     for text in text_list:
         char_list = []
         text = text.translate(god_knows_why_en_testset_contains_zh_quote)
         text = text.translate(custom_trans)
         for seg in jieba.cut(text):
-            seg_byte_len = len(bytes(seg, 'UTF-8'))
+            seg_byte_len = len(bytes(seg, "UTF-8"))
             if seg_byte_len == len(seg):  # if pure alphabets and symbols
                 if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
                     char_list.append(" ")
@@ -186,413 +171,15 @@ 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 = rearrange(padded_ref_mels, 'b d n -> b n d')
-    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., 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:
-            zh_pause_punc = r"。，、；：？！"
-            ref_text_len = len(prompt_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, prompt_text))
-            gen_text_len = len(gt_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gt_text))
-            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 = rearrange(ref_mel, '1 d n -> d n')
-
-        # 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
 
-def repetition_found(text, length = 2, tolerance = 10):
+
+def repetition_found(text, length=2, tolerance=10):
     pattern_count = defaultdict(int)
     for i in range(len(text) - length + 1):
-        pattern = text[i:i + length]
+        pattern = text[i : i + length]
         pattern_count[pattern] += 1
     for pattern, count in pattern_count.items():
         if count > tolerance:
             return True
     return False
-
-
-# load model checkpoint for inference
-
-def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
-    if dtype is None:
-        dtype = (
-            torch.float16 if "cuda" in device and torch.cuda.get_device_properties(device).major >= 6 else torch.float32
-        )
-    model = model.to(dtype)
-
-    ckpt_type = ckpt_path.split(".")[-1]
-    if ckpt_type == "safetensors":
-        from safetensors.torch import load_file
-
-        checkpoint = load_file(ckpt_path, device=device)
-    else:
-        checkpoint = torch.load(ckpt_path, map_location=device, 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"]
-        }
-
-        # patch for backward compatibility, 305e3ea
-        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
-            if key in checkpoint["model_state_dict"]:
-                del checkpoint["model_state_dict"][key]
-
-        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"])
-
-    del checkpoint
-    torch.cuda.empty_cache()
-
-    return model.to(device)